From 827e7c189b08c881762807bc597c4a8a7e166d74 Mon Sep 17 00:00:00 2001 From: Zhichun Wu Date: Thu, 6 May 2021 10:05:08 +0800 Subject: [PATCH 1/2] Ignore database only for DDL containing DATABASE keyword, and allow to change current database of a statement --- .../clickhouse/ClickHouseStatementImpl.java | 13 +++-- .../ClickHouseConnectionImplTest.java | 50 ++++++++++++++++++- 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/ClickHouseStatementImpl.java b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/ClickHouseStatementImpl.java index 7699f5833..47d25c406 100644 --- a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/ClickHouseStatementImpl.java +++ b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/ClickHouseStatementImpl.java @@ -43,6 +43,7 @@ import ru.yandex.clickhouse.except.ClickHouseExceptionSpecifier; import ru.yandex.clickhouse.jdbc.parser.ClickHouseSqlParser; import ru.yandex.clickhouse.jdbc.parser.ClickHouseSqlStatement; +import ru.yandex.clickhouse.jdbc.parser.StatementType; import ru.yandex.clickhouse.response.ClickHouseLZ4Stream; import ru.yandex.clickhouse.response.ClickHouseResponse; import ru.yandex.clickhouse.response.ClickHouseResponseSummary; @@ -141,7 +142,7 @@ public boolean isStreaming() { * between creation of this object and query execution, but javadoc does not allow * {@code setCatalog} influence on already created statements. */ - private final String initialDatabase; + protected String currentDatabase; protected ClickHouseSqlStatement getLastStatement() { ClickHouseSqlStatement stmt = null; @@ -303,7 +304,7 @@ public ClickHouseStatementImpl(CloseableHttpClient client, ClickHouseConnection this.httpContext = ClickHouseHttpClientBuilder.createClientContext(properties); this.connection = connection; this.properties = properties == null ? new ClickHouseProperties() : properties; - this.initialDatabase = this.properties.getDatabase(); + this.currentDatabase = this.properties.getDatabase(); this.isResultSetScrollable = (resultSetType != ResultSet.TYPE_FORWARD_ONLY); this.batchStmts = new ArrayList<>(); @@ -709,7 +710,11 @@ private InputStream getInputStream( Map additionalRequestParams ) throws ClickHouseException { String sql = parsedStmt.getSQL(); - boolean ignoreDatabase = parsedStmt.isRecognized() && !parsedStmt.isDML(); + boolean ignoreDatabase = parsedStmt.isRecognized() && !parsedStmt.isDML() + && parsedStmt.containsKeyword("DATABASE"); + if (parsedStmt.getStatementType() == StatementType.USE) { + currentDatabase = parsedStmt.getDatabaseOrDefault(currentDatabase); + } log.debug("Executing SQL: {}", sql); @@ -866,7 +871,7 @@ private List getUrlQueryParams( Map params = properties.buildQueryParams(true); if (!ignoreDatabase) { - params.put(ClickHouseQueryParam.DATABASE, initialDatabase); + params.put(ClickHouseQueryParam.DATABASE, currentDatabase); } params.putAll(getAdditionalDBParams()); diff --git a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java index 99a7340c8..94c360d4f 100644 --- a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java +++ b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java @@ -13,8 +13,8 @@ import ru.yandex.clickhouse.except.ClickHouseException; import ru.yandex.clickhouse.settings.ClickHouseProperties; +import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; @@ -75,6 +75,54 @@ public void testOofWrongPassword() throws Exception { assertFailure(createDataSource("oof", "baz")); } + @Test + public void testDefaultDatabase() throws Exception { + ClickHouseDataSource ds = ClickHouseContainerForTest.newDataSource(); + String currentDbQuery = "select currentDatabase()"; + try (Connection conn = ds.getConnection(); Statement s = conn.createStatement()) { + try (ResultSet rs = s.executeQuery(currentDbQuery)) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "default"); + assertFalse(rs.next()); + } + s.execute("create database if not exists tdb1; create database if not exists tdb2"); + } + + ds = ClickHouseContainerForTest.newDataSource("tdb2"); + try (Connection conn = ds.getConnection(); Statement s = conn.createStatement()) { + try (ResultSet rs = s.executeQuery("select currentDatabase()")) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "tdb2"); + assertFalse(rs.next()); + } + + s.execute("create table tdb2_aaa(a String) engine=Memory; insert into tdb2_aaa values('3')"); + + try (ResultSet rs = s.executeQuery("select currentDatabase(), a from tdb2_aaa")) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "tdb2"); + assertEquals(rs.getString(2), "3"); + assertFalse(rs.next()); + } + + s.execute("use tdb1; create table tdb1_aaa(a String) engine=Memory; insert into tdb1_aaa values('1')"); + + try (ResultSet rs = s.executeQuery("select currentDatabase(), a from tdb1_aaa")) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "tdb1"); + assertEquals(rs.getString(2), "1"); + assertFalse(rs.next()); + } + + try (ResultSet rs = s.executeQuery("use `tdb2`; select currentDatabase(), a from tdb2_aaa")) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "tdb2"); + assertEquals(rs.getString(2), "3"); + assertFalse(rs.next()); + } + } + } + private static void assertSuccess(DataSource dataSource) throws Exception { Connection connection = dataSource.getConnection(); assertTrue(connection.createStatement().execute("SELECT 1")); From 8d0653426ba8fa999eb4a3b0fb586dab53a1e214 Mon Sep 17 00:00:00 2001 From: Zhichun Wu Date: Thu, 6 May 2021 10:12:39 +0800 Subject: [PATCH 2/2] Enhance unit test to cover prepared statement as well --- .../ClickHouseConnectionImplTest.java | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java index 94c360d4f..d1aa4e3a7 100644 --- a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java +++ b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/integration/ClickHouseConnectionImplTest.java @@ -1,6 +1,7 @@ package ru.yandex.clickhouse.integration; import java.sql.Connection; +import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.Statement; @@ -85,12 +86,19 @@ public void testDefaultDatabase() throws Exception { assertEquals(rs.getString(1), "default"); assertFalse(rs.next()); } + + PreparedStatement p = conn.prepareStatement(currentDbQuery); + try (ResultSet rs = p.executeQuery()) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "default"); + assertFalse(rs.next()); + } s.execute("create database if not exists tdb1; create database if not exists tdb2"); } ds = ClickHouseContainerForTest.newDataSource("tdb2"); try (Connection conn = ds.getConnection(); Statement s = conn.createStatement()) { - try (ResultSet rs = s.executeQuery("select currentDatabase()")) { + try (ResultSet rs = s.executeQuery(currentDbQuery)) { assertTrue(rs.next()); assertEquals(rs.getString(1), "tdb2"); assertFalse(rs.next()); @@ -120,6 +128,14 @@ public void testDefaultDatabase() throws Exception { assertEquals(rs.getString(2), "3"); assertFalse(rs.next()); } + + String sql = "select currentDatabase(), a from tdb2_aaa"; + try (PreparedStatement p = conn.prepareStatement(sql); ResultSet rs = p.executeQuery()) { + assertTrue(rs.next()); + assertEquals(rs.getString(1), "tdb2"); + assertEquals(rs.getString(2), "3"); + assertFalse(rs.next()); + } } }