Add initial namespace to open session

This PR adds initial catalog and schema parameters to `connect`.

Unit tests:
* Passthrough tests
* Test canUseMultipleCatalogs
* Test error behaviour when incorrect namespsace returned
* Test protocol version checks vs initial namespsace

Author: NiallEgan

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

@@ -25,6 +25,8 @@ def __init__(self,
                  access_token: str,
                  metadata: Optional[List[Tuple[str, str]]] = None,
                  session_configuration: Dict[str, Any] = None,
+                 catalog: Optional[str] = None,
+                 schema: Optional[str] = None,
                  **kwargs) -> None:
         """
         Connect to a Databricks SQL endpoint or a Databricks cluster.
@@ -36,6 +38,8 @@ def __init__(self,
         :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.
+        :param catalog: An optional initial catalog to use. Requires DBR version 9.0+
+        :param schema: An optional initial schema to use. Requires DBR version 9.0+
         """
 
         # Internal arguments in **kwargs:
@@ -88,7 +92,8 @@ 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(session_configuration)
+        self._session_handle = self.thrift_backend.open_session(session_configuration, catalog,
+                                                                schema)
         self.open = True
         logger.info("Successfully opened session " + str(self.get_session_id()))
         self._cursors = []  # type: List[Cursor]

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

@@ -1,3 +1,4 @@
+### PEP-249 Mandated ###
 class Error(Exception):
     pass
 
@@ -36,3 +37,9 @@ class DataError(DatabaseError):
 
 class NotSupportedError(DatabaseError):
     pass
+
+
+### Custom error classes ###
+class InvalidServerResponseError(OperationalError):
+    """ Thrown if the server does not set the initial namespace correctly"""
+    pass

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

@@ -289,6 +289,22 @@ def _check_protocol_version(self, t_open_session_resp):
                                    "SPARK_CLI_SERVICE_PROTOCOL_V3, "
                                    "instead got: {}".format(protocol_version))
 
+    def _check_initial_namespace(self, catalog, schema, response):
+        if not (catalog or schema):
+            return
+
+        if response.serverProtocolVersion < \
+                ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4:
+            raise InvalidServerResponseError(
+                "Setting initial namespace not supported by the DBR version, "
+                "Please use a Databricks SQL endpoint or a cluster with DBR >= 9.0.")
+
+        if catalog:
+            if not response.canUseMultipleCatalogs:
+                raise InvalidServerResponseError(
+                    "Unexpected response from server: Trying to set initial catalog to {}, " +
+                    "but server does not support multiple catalogs.".format(catalog))
+
     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":
@@ -297,7 +313,7 @@ def _check_session_configuration(self, session_configuration):
                             TIMESTAMP_AS_STRING_CONFIG,
                             session_configuration[TIMESTAMP_AS_STRING_CONFIG]))
 
-    def open_session(self, session_configuration):
+    def open_session(self, session_configuration, catalog, schema):
         try:
             self._transport.open()
             session_configuration = {k: str(v) for (k, v) in (session_configuration or {}).items()}
@@ -306,11 +322,19 @@ def open_session(self, session_configuration):
             # 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"
+            if catalog or schema:
+                initial_namespace = ttypes.TNamespace(catalogName=catalog, schemaName=schema)
+            else:
+                initial_namespace = None
+
             open_session_req = ttypes.TOpenSessionReq(
                 client_protocol_i64=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
                 client_protocol=None,
+                initialNamespace=initial_namespace,
+                canUseMultipleCatalogs=True,
                 configuration=session_configuration)
             response = self.make_request(self._client.OpenSession, open_session_req)
+            self._check_initial_namespace(catalog, schema, response)
             self._check_protocol_version(response)
             return response.sessionHandle
         except:

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({}, None, None)
 
             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({}, None, None)
 
     @patch("thrift.transport.THttpClient.THttpClient")
     def test_headers_are_set(self, t_http_client_class):
@@ -673,7 +673,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({}, None, None)
         self.assertEqual(len(tcli_service_instance.OpenSession.call_args_list), 1)
 
     @patch("databricks.sql.thrift_backend.TCLIService.Client")
@@ -1057,7 +1057,7 @@ def test_configuration_passthrough(self, tcli_client_class):
         }
 
         backend = ThriftBackend("foobar", 443, "path", [])
-        backend.open_session(mock_config)
+        backend.open_session(mock_config, None, None)
 
         open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
         self.assertEqual(open_session_req.configuration, expected_config)
@@ -1070,10 +1070,92 @@ def test_cant_set_timestamp_as_string_to_true(self, tcli_client_class):
         backend = ThriftBackend("foobar", 443, "path", [])
 
         with self.assertRaises(databricks.sql.Error) as cm:
-            backend.open_session(mock_config)
+            backend.open_session(mock_config, None, None)
 
         self.assertIn("timestampAsString cannot be changed", str(cm.exception))
 
+    def _construct_open_session_with_namespace(self, can_use_multiple_cats, cat, schem):
+        return ttypes.TOpenSessionResp(
+            status=self.okay_status,
+            serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
+            canUseMultipleCatalogs=can_use_multiple_cats,
+            initialNamespace=ttypes.TNamespace(catalogName=cat, schemaName=schem))
+
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    def test_initial_namespace_passthrough_to_open_session(self, tcli_client_class):
+        tcli_service_instance = tcli_client_class.return_value
+
+        backend = ThriftBackend("foobar", 443, "path", [])
+        initial_cat_schem_args = [("cat", None), (None, "schem"), ("cat", "schem")]
+
+        for cat, schem in initial_cat_schem_args:
+            with self.subTest(cat=cat, schem=schem):
+                tcli_service_instance.OpenSession.return_value = \
+                    self._construct_open_session_with_namespace(True, cat, schem)
+
+                backend.open_session({}, cat, schem)
+
+                open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
+                self.assertEqual(open_session_req.initialNamespace.catalogName, cat)
+                self.assertEqual(open_session_req.initialNamespace.schemaName, schem)
+
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    def test_can_use_multiple_catalogs_is_set_in_open_session_req(self, tcli_client_class):
+        tcli_service_instance = tcli_client_class.return_value
+        tcli_service_instance.OpenSession.return_value = self.open_session_resp
+
+        backend = ThriftBackend("foobar", 443, "path", [])
+        backend.open_session({}, None, None)
+
+        open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
+        self.assertTrue(open_session_req.canUseMultipleCatalogs)
+
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    def test_can_use_multiple_catalogs_is_false_fails_with_initial_catalog(self, tcli_client_class):
+        tcli_service_instance = tcli_client_class.return_value
+
+        backend = ThriftBackend("foobar", 443, "path", [])
+        # If the initial catalog is set, but server returns canUseMultipleCatalogs=False, we
+        # expect failure. If the initial catalog isn't set, then canUseMultipleCatalogs=False
+        # is fine
+        failing_ns_args = [("cat", None), ("cat", "schem")]
+        passing_ns_args = [(None, None), (None, "schem")]
+
+        for cat, schem in failing_ns_args:
+            tcli_service_instance.OpenSession.return_value = \
+                self._construct_open_session_with_namespace(False, cat, schem)
+
+            with self.assertRaises(InvalidServerResponseError) as cm:
+                backend.open_session({}, cat, schem)
+
+            self.assertIn("server does not support multiple catalogs", str(cm.exception),
+                          "incorrect error thrown for initial namespace {}".format((cat, schem)))
+
+        for cat, schem in passing_ns_args:
+            tcli_service_instance.OpenSession.return_value = \
+                self._construct_open_session_with_namespace(False, cat, schem)
+            backend.open_session({}, cat, schem)
+
+    @patch("databricks.sql.thrift_backend.TCLIService.Client")
+    def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
+        tcli_service_instance = tcli_client_class.return_value
+
+        tcli_service_instance.OpenSession.return_value = \
+            ttypes.TOpenSessionResp(
+                status=self.okay_status,
+                serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V3,
+                canUseMultipleCatalogs=True,
+                initialNamespace=ttypes.TNamespace(catalogName="cat", schemaName="schem")
+            )
+
+        backend = ThriftBackend("foobar", 443, "path", [])
+
+        with self.assertRaises(InvalidServerResponseError) as cm:
+            backend.open_session({}, "cat", "schem")
+
+        self.assertIn("Setting initial namespace not supported by the DBR version",
+                      str(cm.exception))
+
 
 if __name__ == '__main__':
     unittest.main()

cmdexec/clients/python/tests/tests.py

@@ -353,7 +353,16 @@ def test_configuration_passthrough(self, mock_client_class):
         self.assertEqual(mock_client_class.return_value.open_session.call_args[0][0],
                          mock_session_config)
 
-    def test_execute_parameter_passhthrough(self):
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_initial_namespace_passthrough(self, mock_client_class):
+        mock_cat = Mock()
+        mock_schem = Mock()
+
+        databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, catalog=mock_cat, schema=mock_schem)
+        self.assertEqual(mock_client_class.return_value.open_session.call_args[0][1], mock_cat)
+        self.assertEqual(mock_client_class.return_value.open_session.call_args[0][2], mock_schem)
+
+    def test_execute_parameter_passthrough(self):
         mock_thrift_backend = Mock()
         cursor = client.Cursor(Mock(), mock_thrift_backend)