Merge pull request #312 from databendlabs/feat/support-session-settings-pros

feat: support session settings in properties

Author: hantmac

databend-client/src/main/java/com/databend/client/ClientSettings.java

@@ -44,6 +44,7 @@ public class ClientSettings {
     private final int retryAttempts;
     // TODO(zhihanz) timezone and locale info
 
+    //ClientSettings for test case use
     public ClientSettings(String host) {
         this(host, DatabendSession.createDefault(), DEFAULT_QUERY_TIMEOUT, DEFAULT_CONNECTION_TIMEOUT, DEFAULT_SOCKET_TIMEOUT, PaginationOptions.defaultPaginationOptions(), new HashMap<String, String>(), null, DEFAULT_RETRY_ATTEMPTS);
     }

databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java

@@ -43,6 +43,8 @@ public final class ConnectionProperties {
     public static final ConnectionProperty<Integer> MAX_ROWS_IN_BUFFER = new MaxRowsInBuffer();
     public static final ConnectionProperty<Integer> MAX_ROWS_PER_PAGE = new MaxRowsPerPage();
 
+    public static final ConnectionProperty<String> SESSION_SETTINGS = new SessionSettings();
+
     private static final Set<ConnectionProperty<?>> ALL_PROPERTIES = ImmutableSet.<ConnectionProperty<?>>builder()
             .add(USER)
             .add(PASSWORD)
@@ -63,6 +65,7 @@ public final class ConnectionProperties {
             .add(WAIT_TIME_SECS)
             .add(MAX_ROWS_IN_BUFFER)
             .add(MAX_ROWS_PER_PAGE)
+            .add(SESSION_SETTINGS)
             .build();
     private static final Map<String, String> DEFAULTS;
 
@@ -250,6 +253,13 @@ public MaxRowsPerPage() {
         }
     }
 
+    private static class SessionSettings
+            extends AbstractConnectionProperty<String> {
+        public SessionSettings() {
+            super("session_settings", Optional.of(""), NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
+        }
+    }
+
 
     static {
         ImmutableMap.Builder<String, String> defaults = ImmutableMap.builder();

databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java

@@ -124,7 +124,7 @@ private void initializeFileHandler() {
         this.routeHint = randRouteHint();
         // it maybe closed due to unsupported server versioning.
         this.autoDiscovery = uri.autoDiscovery();
-        DatabendSession session = new DatabendSession.Builder().setDatabase(this.getSchema()).build();
+        DatabendSession session = new DatabendSession.Builder().setDatabase(this.getSchema()).setSettings(uri.getSessionSettings()).build();
         this.setSession(session);
 
         initializeFileHandler();

databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java

@@ -67,6 +67,8 @@ public final class DatabendDriverUri {
     private final Integer maxRowsPerPage;
     private final int nodeDiscoveryInterval;
 
+    private final Map<String, String> sessionSettings;
+
 //    private final boolean useSecureConnection;
 
     private DatabendDriverUri(String url, Properties driverProperties)
@@ -104,8 +106,28 @@ private DatabendDriverUri(String url, Properties driverProperties)
         this.socketTimeout = SOCKET_TIMEOUT.getRequiredValue(properties);
         this.maxRowsInBuffer = ConnectionProperties.MAX_ROWS_IN_BUFFER.getRequiredValue(properties);
         this.maxRowsPerPage = ConnectionProperties.MAX_ROWS_PER_PAGE.getRequiredValue(properties);
+        String settingsStr = SESSION_SETTINGS.getValue(properties).orElse("");
+        this.sessionSettings = parseSessionSettings(settingsStr);
+    }
+
+    private Map<String, String> parseSessionSettings(String settingsStr) {
+        if (isNullOrEmpty(settingsStr)) {
+            return new HashMap<>();
+        }
+
+        // key1=value1,key2=value2
+        Map<String, String> settings = new HashMap<>();
+        String[] pairs = settingsStr.split(",");
+        for (String pair : pairs) {
+            String[] keyValue = pair.split("=", 2);
+            if (keyValue.length == 2) {
+                settings.put(keyValue[0].trim(), keyValue[1].trim());
+            }
+        }
+        return settings;
     }
 
+
     public static DatabendDriverUri create(String url, Properties properties)
             throws SQLException {
         return new DatabendDriverUri(url, firstNonNull(properties, new Properties()));
@@ -406,6 +428,10 @@ public Integer getMaxFailoverRetry() {
         return maxFailoverRetry;
     }
 
+    public Map<String, String> getSessionSettings() {
+        return sessionSettings;
+    }
+
     public Properties getProperties() {
         return properties;
     }

databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java

@@ -4,7 +4,9 @@
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.Properties;
 
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -246,4 +248,32 @@ public void TestSetSchema() throws SQLException {
         Assert.assertEquals(connection.getSchema(), "test2");
         connection.createStatement().execute("insert into test2 values (1)");
     }
+
+    @Test
+    public void TestSetSessionSettings() throws SQLException{
+        Properties props = new Properties();
+        // set session settings
+        props.setProperty("session_settings", "key1=value1,key2=value2");
+        props.setProperty("user","databend");
+        props.setProperty("password","databend");
+        DatabendConnection connection = (DatabendConnection) Utils.createConnection("default",props);
+        try {
+            Statement statement = connection.createStatement();
+            statement.execute("show settings");
+            ResultSet r = statement.getResultSet();
+            while (r.next()) {
+                String name = r.getString("name");
+                String value = r.getString("value");
+                if (name.equals("key1")) {
+                    Assert.assertEquals(value, "value1");
+                } else if (name.equals("key2")) {
+                    Assert.assertEquals(value, "value2");
+                }
+            }
+        } catch (SQLException e) {
+            throw new RuntimeException(e);
+        } finally {
+            connection.close();
+        }
+    }
 }