Fix issue #648 introduced in 0.3.1

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<String, String> 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<NameValuePair> getUrlQueryParams(
 
         Map<ClickHouseQueryParam, String> params = properties.buildQueryParams(true);
         if (!ignoreDatabase) {
-            params.put(ClickHouseQueryParam.DATABASE, initialDatabase);
+            params.put(ClickHouseQueryParam.DATABASE, currentDatabase);
         }
 
         params.putAll(getAdditionalDBParams());

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;
 
@@ -13,8 +14,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 +76,69 @@ 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());
+            }
+
+            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(currentDbQuery)) {
+                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());
+            }
+
+            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());
+            }
+        }
+    }
+
     private static void assertSuccess(DataSource dataSource) throws Exception {
         Connection connection = dataSource.getConnection();
         assertTrue(connection.createStatement().execute("SELECT 1"));