Add session state properties passthrough

This PR adds the `session_configuration` parameter to allow users to set confs on the session.

Author: NiallEgan

cmdexec/clients/python/src/databricks/sql/client.py

@@ -24,6 +24,7 @@ def __init__(self,
                  http_path: str,
                  access_token: str,
                  metadata: Optional[List[Tuple[str, str]]] = None,
+                 session_configuration: Dict[str, Any] = None,
                  **kwargs) -> None:
         """
         Connect to a Databricks SQL endpoint or a Databricks cluster.
@@ -33,6 +34,8 @@ def __init__(self,
               or to a DBR interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)
         :param access_token: Http Bearer access token, e.g. Databricks Personal Access Token.
         :param metadata: An optional list of (k, v) pairs that will be set as Http headers on every request
+        :param session_configuration: An optional dictionary of Spark session parameters. Defaults to None. 
+               Execute the SQL command `SET -v` to get a full list of available commands.
         """
 
         # Internal arguments in **kwargs:
@@ -82,7 +85,7 @@ def __init__(self,
         self.thrift_backend = ThriftBackend(self.host, self.port, http_path,
                                             (metadata or []) + base_headers, **kwargs)
 
-        self._session_handle = self.thrift_backend.open_session()
+        self._session_handle = self.thrift_backend.open_session(session_configuration)
         self.open = True
         logger.info("Successfully opened session " + str(self.get_session_id()))
         self._cursors = []  # type: List[Cursor]

cmdexec/clients/python/src/databricks/sql/thrift_backend.py

@@ -23,6 +23,8 @@
 DATABRICKS_ERROR_OR_REDIRECT_HEADER = "x-databricks-error-or-redirect-message"
 DATABRICKS_REASON_HEADER = "x-databricks-reason-phrase"
 
+TIMESTAMP_AS_STRING_CONFIG = "spark.thriftserver.arrowBasedRowSet.timestampAsString"
+
 # see Connection.__init__ for parameter descriptions.
 # - Min/Max avoids unsustainable configs (sane values are far more constrained)
 # - 900s attempts-duration lines up w ODBC/JDBC drivers (for cluster startup > 10 mins)
@@ -278,18 +280,27 @@ def _check_protocol_version(self, t_open_session_resp):
                                    "SPARK_CLI_SERVICE_PROTOCOL_V3, "
                                    "instead got: {}".format(protocol_version))
 
-    def open_session(self):
+    def _check_session_configuration(self, session_configuration):
+        # This client expects timetampsAsString to be false, so we do not allow users to modify that
+        if session_configuration.get(TIMESTAMP_AS_STRING_CONFIG, "false").lower() != "false":
+            raise Error("Invalid session configuration: {} cannot be changed "
+                        "while using the Databricks SQL connector, it must be false not {}".format(
+                            TIMESTAMP_AS_STRING_CONFIG,
+                            session_configuration[TIMESTAMP_AS_STRING_CONFIG]))
+
+    def open_session(self, session_configuration):
         try:
             self._transport.open()
+            session_configuration = {k: str(v) for (k, v) in (session_configuration or {}).items()}
+            self._check_session_configuration(session_configuration)
+            # We want to receive proper Timestamp arrow types.
+            # We set it also in confOverlay in TExecuteStatementReq on a per query basic,
+            # but it doesn't hurt to also set for the whole session.
+            session_configuration[TIMESTAMP_AS_STRING_CONFIG] = "false"
             open_session_req = ttypes.TOpenSessionReq(
                 client_protocol_i64=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
                 client_protocol=None,
-                configuration={
-                    # We want to receive proper Timestamp arrow types.
-                    # We set it also in confOverlay in TExecuteStatementReq on a per query basic,
-                    # but it doesn't hurt to also set for the whole session.
-                    "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
-                })
+                configuration=session_configuration)
             response = self.make_request(self._client.OpenSession, open_session_req)
             self._check_protocol_version(response)
             return response.sessionHandle

cmdexec/clients/python/tests/test_thrift_backend.py

@@ -117,7 +117,7 @@ def test_bad_protocol_versions_are_rejected(self, tcli_service_client_cass):
 
             with self.assertRaises(OperationalError) as cm:
                 thrift_backend = self._make_fake_thrift_backend()
-                thrift_backend.open_session()
+                thrift_backend.open_session({})
 
             self.assertIn("expected server to use a protocol version", str(cm.exception))
 
@@ -134,7 +134,7 @@ def test_okay_protocol_versions_succeed(self, tcli_service_client_cass):
                 status=self.okay_status, serverProtocolVersion=protocol_version)
 
             thrift_backend = self._make_fake_thrift_backend()
-            thrift_backend.open_session()
+            thrift_backend.open_session({})
 
     @patch("thrift.transport.THttpClient.THttpClient")
     def test_headers_are_set(self, t_http_client_class):
@@ -664,7 +664,7 @@ def test_open_session_user_provided_session_id_optional(self, tcli_service_class
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
 
         thrift_backend = ThriftBackend("foobar", 443, "path", [])
-        thrift_backend.open_session()
+        thrift_backend.open_session({})
         self.assertEqual(len(tcli_service_instance.OpenSession.call_args_list), 1)
 
     @patch("databricks.sql.thrift_backend.TCLIService.Client")
@@ -1035,6 +1035,36 @@ def test_retry_args_bounding(self, mock_http_client):
             for (arg, val) in retry_delay_expected_vals.items():
                 self.assertEqual(getattr(backend, arg), val)
 
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    def test_configuration_passthrough(self, tcli_client_class):
+        tcli_service_instance = tcli_client_class.return_value
+        tcli_service_instance.OpenSession.return_value = self.open_session_resp
+        mock_config = {"foo": "bar", "baz": True, "42": 42}
+        expected_config = {
+            "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false",
+            "foo": "bar",
+            "baz": "True",
+            "42": "42"
+        }
+
+        backend = ThriftBackend("foobar", 443, "path", [])
+        backend.open_session(mock_config)
+
+        open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
+        self.assertEqual(open_session_req.configuration, expected_config)
+
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    def test_cant_set_timestamp_as_string_to_true(self, tcli_client_class):
+        tcli_service_instance = tcli_client_class.return_value
+        tcli_service_instance.OpenSession.return_value = self.open_session_resp
+        mock_config = {"spark.thriftserver.arrowBasedRowSet.timestampAsString": True}
+        backend = ThriftBackend("foobar", 443, "path", [])
+
+        with self.assertRaises(databricks.sql.Error) as cm:
+            backend.open_session(mock_config)
+
+        self.assertIn("timestampAsString cannot be changed", str(cm.exception))
+
 
 if __name__ == '__main__':
     unittest.main()

cmdexec/clients/python/tests/tests.py

@@ -339,6 +339,15 @@ def test_version_is_canonical(self):
                                r'(0|[1-9][0-9]*))?(\.post(0|[1-9][0-9]*))?(\.dev(0|[1-9][0-9]*))?$'
         self.assertIsNotNone(re.match(canonical_version_re, version))
 
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_configuration_passthrough(self, mock_client_class):
+        mock_session_config = Mock()
+        databricks.sql.connect(
+            session_configuration=mock_session_config, **self.DUMMY_CONNECTION_ARGS)
+
+        self.assertEqual(mock_client_class.return_value.open_session.call_args[0][0],
+                         mock_session_config)
+
 
 if __name__ == '__main__':
     suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])