[client-v2] Session API with reusable Session objects

client-v2/src/main/java/com/clickhouse/client/api/Client.java

@@ -117,6 +117,7 @@ public class Client implements AutoCloseable {
 
     private final List<Endpoint> endpoints;
     private final Map<String, Object> configuration;
+    private final Session session;
 
     private final Map<String, String> readOnlyConfig;
 
@@ -145,7 +146,9 @@ public class Client implements AutoCloseable {
     private Client(Collection<Endpoint> endpoints, Map<String,String> configuration,
                    ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy,
                    Object metricsRegistry, Supplier<String> queryIdGenerator) {
-        this.configuration = ClientConfigProperties.parseConfigMap(configuration);
+        Map<String, Object> parsedConfiguration = ClientConfigProperties.parseConfigMap(configuration);
+        this.session = Session.extractFrom(parsedConfiguration);
+        this.configuration = new ConcurrentHashMap<>(parsedConfiguration);
         this.readOnlyConfig = Collections.unmodifiableMap(configuration);
         this.metricsRegistry = metricsRegistry;
         this.queryIdGenerator = queryIdGenerator;
@@ -955,6 +958,54 @@ public Builder serverSetting(String name,  Collection<String> values) {
             return this;
         }
 
+        /**
+         * Sets ClickHouse session id to be sent with each request.
+         */
+        public Builder setSessionId(String sessionId) {
+            ValidationUtils.checkNonBlank(sessionId, ClickHouseHttpProto.QPARAM_SESSION_ID);
+            return serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID, sessionId);
+        }
+
+        /**
+         * Sets ClickHouse session check flag to be sent with each request.
+         */
+        public Builder setSessionCheck(boolean sessionCheck) {
+            return serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK, sessionCheck ? "1" : "0");
+        }
+
+        /**
+         * Sets ClickHouse session timeout in seconds to be sent with each request.
+         */
+        public Builder setSessionTimeout(int timeoutInSeconds) {
+            ValidationUtils.checkPositive(timeoutInSeconds, ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT);
+            return serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT, String.valueOf(timeoutInSeconds));
+        }
+
+        /**
+         * Sets ClickHouse session timezone to be sent with each request.
+         */
+        public Builder setSessionTimezone(String timezone) {
+            ValidationUtils.checkNonBlank(timezone, ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE);
+            return serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE, timezone);
+        }
+
+        public Builder use(Session session) {
+            ValidationUtils.checkNotNull(session, "session");
+            if (session.getSessionId() != null) {
+                setSessionId(session.getSessionId());
+            }
+            if (session.getSessionCheck() != null) {
+                setSessionCheck(session.getSessionCheck());
+            }
+            if (session.getSessionTimeout() != null) {
+                setSessionTimeout(session.getSessionTimeout());
+            }
+            if (session.getSessionTimezone() != null) {
+                setSessionTimezone(session.getSessionTimezone());
+            }
+            return this;
+        }
+
         /**
          * Sets column to method matching strategy. It is used while registering POJO serializers and deserializers.
          * Default is {@link DefaultColumnToMethodMatchingStrategy}.
@@ -2138,6 +2189,14 @@ public void updateClientName(String name) {
         this.configuration.put(ClientConfigProperties.CLIENT_NAME.getKey(), name);
     }
 
+    /**
+     * Updates ClickHouse session id for all subsequent requests created by this client.
+     */
+    public void updateSessionId(String sessionId) {
+        ValidationUtils.checkNonBlank(sessionId, ClickHouseHttpProto.QPARAM_SESSION_ID);
+        this.session.updateSessionId(sessionId);
+    }
+
     public static final String clientVersion =
             ClickHouseClientOption.readVersionFromResource("client-v2-version.properties");
     public static final String CLIENT_USER_AGENT = "clickhouse-java-v2/";
@@ -2170,8 +2229,8 @@ private Endpoint getNextAliveNode() {
      * @return request settings - merged client and operation settings
      */
     private Map<String, Object> buildRequestSettings(Map<String, Object> opSettings) {
-        Map<String, Object> requestSettings = new HashMap<>();
-        requestSettings.putAll(configuration);
+        Map<String, Object> requestSettings = new HashMap<>(configuration);
+        session.applyTo(requestSettings);
         requestSettings.putAll(opSettings);
         return requestSettings;
     }

client-v2/src/main/java/com/clickhouse/client/api/Session.java

@@ -0,0 +1,100 @@
+package com.clickhouse.client.api;
+
+import com.clickhouse.client.api.http.ClickHouseHttpProto;
+import com.clickhouse.client.api.internal.ValidationUtils;
+
+import java.util.Map;
+
+/**
+ * Reusable ClickHouse session configuration that can be applied to clients or operation settings.
+ */
+public class Session {
+    private String sessionId;
+    private Boolean sessionCheck;
+    private Integer sessionTimeout;
+    private String sessionTimezone;
+
+    static Session extractFrom(Map<String, Object> configuration) {
+        Session session = new Session();
+
+        String sessionId = (String) configuration.remove(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID));
+        if (sessionId != null) {
+            session.setSessionId(sessionId);
+        }
+
+        String sessionCheck = (String) configuration.remove(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK));
+        if (sessionCheck != null) {
+            session.setSessionCheck("1".equals(sessionCheck) || Boolean.parseBoolean(sessionCheck));
+        }
+
+        String sessionTimeout = (String) configuration.remove(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT));
+        if (sessionTimeout != null) {
+            session.setSessionTimeout(Integer.parseInt(sessionTimeout));
+        }
+
+        String sessionTimezone = (String) configuration.remove(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE));
+        if (sessionTimezone != null) {
+            session.setSessionTimezone(sessionTimezone);
+        }
+
+        return session;
+    }
+
+    public synchronized Session setSessionId(String sessionId) {
+        ValidationUtils.checkNonBlank(sessionId, ClickHouseHttpProto.QPARAM_SESSION_ID);
+        this.sessionId = sessionId;
+        return this;
+    }
+
+    public synchronized String getSessionId() {
+        return sessionId;
+    }
+
+    public synchronized Session setSessionCheck(boolean sessionCheck) {
+        this.sessionCheck = sessionCheck;
+        return this;
+    }
+
+    public synchronized Boolean getSessionCheck() {
+        return sessionCheck;
+    }
+
+    public synchronized Session setSessionTimeout(int timeoutInSeconds) {
+        ValidationUtils.checkPositive(timeoutInSeconds, ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT);
+        this.sessionTimeout = timeoutInSeconds;
+        return this;
+    }
+
+    public synchronized Integer getSessionTimeout() {
+        return sessionTimeout;
+    }
+
+    public synchronized Session setSessionTimezone(String timezone) {
+        ValidationUtils.checkNonBlank(timezone, ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE);
+        this.sessionTimezone = timezone;
+        return this;
+    }
+
+    public synchronized String getSessionTimezone() {
+        return sessionTimezone;
+    }
+
+    public synchronized void updateSessionId(String sessionId) {
+        setSessionId(sessionId);
+    }
+
+    public synchronized void applyTo(Map<String, Object> requestSettings) {
+        putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_ID, sessionId);
+        putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_CHECK,
+                sessionCheck == null ? null : (sessionCheck ? "1" : "0"));
+        putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT,
+                sessionTimeout == null ? null : String.valueOf(sessionTimeout));
+        putIfSet(requestSettings, ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE, sessionTimezone);
+    }
+
+    private static void putIfSet(Map<String, Object> settings, String key, String value) {
+        if (value != null) {
+            settings.put(ClientConfigProperties.serverSetting(key), value);
+        }
+    }
+}

client-v2/src/main/java/com/clickhouse/client/api/command/CommandSettings.java

@@ -1,6 +1,36 @@
 package com.clickhouse.client.api.command;
 
+import com.clickhouse.client.api.Session;
 import com.clickhouse.client.api.query.QuerySettings;
 
 public class CommandSettings extends QuerySettings {
+    @Override
+    public CommandSettings setSessionId(String sessionId) {
+        super.setSessionId(sessionId);
+        return this;
+    }
+
+    @Override
+    public CommandSettings setSessionCheck(boolean sessionCheck) {
+        super.setSessionCheck(sessionCheck);
+        return this;
+    }
+
+    @Override
+    public CommandSettings setSessionTimeout(int timeoutInSeconds) {
+        super.setSessionTimeout(timeoutInSeconds);
+        return this;
+    }
+
+    @Override
+    public CommandSettings setSessionTimezone(String timezone) {
+        super.setSessionTimezone(timezone);
+        return this;
+    }
+
+    @Override
+    public CommandSettings use(Session session) {
+        super.use(session);
+        return this;
+    }
 }

client-v2/src/main/java/com/clickhouse/client/api/http/ClickHouseHttpProto.java

@@ -57,6 +57,26 @@ public class ClickHouseHttpProto {
      */
     public static final String QPARAM_QUERY_ID = "query_id";
 
+    /**
+     * Query parameter to specify a session id.
+     */
+    public static final String QPARAM_SESSION_ID = "session_id";
+
+    /**
+     * Query parameter to check session status (1/0).
+     */
+    public static final String QPARAM_SESSION_CHECK = "session_check";
+
+    /**
+     * Query parameter to specify session timeout in seconds.
+     */
+    public static final String QPARAM_SESSION_TIMEOUT = "session_timeout";
+
+    /**
+     * Query parameter to specify session timezone.
+     */
+    public static final String QPARAM_SESSION_TIMEZONE = "session_timezone";
+
     public static final String QPARAM_ROLE = "role";
 
     /**

client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java

@@ -2,6 +2,7 @@
 
 import com.clickhouse.client.api.Client;
 import com.clickhouse.client.api.ClientConfigProperties;
+import com.clickhouse.client.api.Session;
 import com.clickhouse.client.api.internal.CommonSettings;
 import org.apache.hc.core5.http.HttpHeaders;
 
@@ -91,6 +92,59 @@ public InsertSettings setQueryId(String queryId) {
         return this;
     }
 
+    /**
+     * Sets ClickHouse session id for this operation.
+     */
+    public InsertSettings setSessionId(String sessionId) {
+        settings.setSessionId(sessionId);
+        return this;
+    }
+
+    public String getSessionId() {
+        return settings.getSessionId();
+    }
+
+    /**
+     * Sets ClickHouse session check flag for this operation.
+     */
+    public InsertSettings setSessionCheck(boolean sessionCheck) {
+        settings.setSessionCheck(sessionCheck);
+        return this;
+    }
+
+    public Boolean getSessionCheck() {
+        return settings.getSessionCheck();
+    }
+
+    /**
+     * Sets ClickHouse session timeout (seconds) for this operation.
+     */
+    public InsertSettings setSessionTimeout(int timeoutInSeconds) {
+        settings.setSessionTimeout(timeoutInSeconds);
+        return this;
+    }
+
+    public Integer getSessionTimeout() {
+        return settings.getSessionTimeout();
+    }
+
+    /**
+     * Sets ClickHouse session timezone for this operation.
+     */
+    public InsertSettings setSessionTimezone(String timezone) {
+        settings.setSessionTimezone(timezone);
+        return this;
+    }
+
+    public String getSessionTimezone() {
+        return settings.getSessionTimezone();
+    }
+
+    public InsertSettings use(Session session) {
+        settings.use(session);
+        return this;
+    }
+
     public int getInputStreamCopyBufferSize() {
         return this.inputStreamCopyBufferSize;
     }

client-v2/src/main/java/com/clickhouse/client/api/internal/CommonSettings.java

@@ -2,6 +2,8 @@
 
 import com.clickhouse.client.api.Client;
 import com.clickhouse.client.api.ClientConfigProperties;
+import com.clickhouse.client.api.Session;
+import com.clickhouse.client.api.http.ClickHouseHttpProto;
 
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
@@ -90,6 +92,53 @@ public CommonSettings setQueryId(String queryId) {
         return this;
     }
 
+    public CommonSettings setSessionId(String sessionId) {
+        ValidationUtils.checkNonBlank(sessionId, ClickHouseHttpProto.QPARAM_SESSION_ID);
+        serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID, sessionId);
+        return this;
+    }
+
+    public String getSessionId() {
+        return (String) settings.get(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_ID));
+    }
+
+    public CommonSettings setSessionCheck(boolean sessionCheck) {
+        serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK, sessionCheck ? "1" : "0");
+        return this;
+    }
+
+    public Boolean getSessionCheck() {
+        String value = (String) settings.get(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_CHECK));
+        return value == null ? null : ("1".equals(value) || Boolean.parseBoolean(value));
+    }
+
+    public CommonSettings setSessionTimeout(int timeoutInSeconds) {
+        ValidationUtils.checkPositive(timeoutInSeconds, ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT);
+        serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT, String.valueOf(timeoutInSeconds));
+        return this;
+    }
+
+    public Integer getSessionTimeout() {
+        String value = (String) settings.get(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEOUT));
+        return value == null ? null : Integer.valueOf(value);
+    }
+
+    public CommonSettings setSessionTimezone(String timezone) {
+        ValidationUtils.checkNonBlank(timezone, ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE);
+        serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE, timezone);
+        return this;
+    }
+
+    public String getSessionTimezone() {
+        return (String) settings.get(ClientConfigProperties.serverSetting(ClickHouseHttpProto.QPARAM_SESSION_TIMEZONE));
+    }
+
+    public CommonSettings use(Session session) {
+        ValidationUtils.checkNotNull(session, "session");
+        session.applyTo(settings);
+        return this;
+    }
+
     /**
      * Operation id. Used internally to register new operation.
      * Should not be called directly.