use factory for backend instantiation

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/session.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, Tuple, List, Optional, Any
+from typing import Dict, Tuple, List, Optional, Any, Type
 
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
@@ -62,6 +62,7 @@ def __init__(
             useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__)
 
         base_headers = [("User-Agent", useragent_header)]
+        all_headers = (http_headers or []) + base_headers
 
         self._ssl_options = SSLOptions(
             # Double negation is generally a bad thing, but we have to keep backward compatibility
@@ -75,33 +76,49 @@ def __init__(
             tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
         )
 
-        # Determine which backend to use
+        self.backend = self._create_backend(
+            server_hostname,
+            http_path,
+            all_headers,
+            auth_provider,
+            _use_arrow_native_complex_types,
+            kwargs,
+        )
+
+        self.protocol_version = None
+
+    def _create_backend(
+        self,
+        server_hostname: str,
+        http_path: str,
+        all_headers: List[Tuple[str, str]],
+        auth_provider,
+        _use_arrow_native_complex_types: bool,
+        kwargs: dict,
+    ) -> DatabricksClient:
+        """Create and return the appropriate backend client."""
         use_sea = kwargs.get("use_sea", False)
 
         if use_sea:
-            self.backend: DatabricksClient = SeaDatabricksClient(
-                self.host,
-                self.port,
-                http_path,
-                (http_headers or []) + base_headers,
-                auth_provider,
-                ssl_options=self._ssl_options,
-                _use_arrow_native_complex_types=_use_arrow_native_complex_types,
-                **kwargs,
-            )
+            logger.debug("Creating SEA backend client")
+            databricks_client_class = SeaDatabricksClient
         else:
-            self.backend = ThriftDatabricksClient(
-                self.host,
-                self.port,
-                http_path,
-                (http_headers or []) + base_headers,
-                auth_provider,
-                ssl_options=self._ssl_options,
-                _use_arrow_native_complex_types=_use_arrow_native_complex_types,
-                **kwargs,
-            )
-
-        self.protocol_version = None
+            logger.debug("Creating Thrift backend client")
+            databricks_client_class = ThriftDatabricksClient
+
+        # Prepare common arguments
+        common_args = {
+            "server_hostname": server_hostname,
+            "port": self.port,
+            "http_path": http_path,
+            "http_headers": all_headers,
+            "auth_provider": auth_provider,
+            "ssl_options": self._ssl_options,
+            "_use_arrow_native_complex_types": _use_arrow_native_complex_types,
+            **kwargs,
+        }
+
+        return databricks_client_class(**common_args)
 
     def open(self):
         self._session_id = self.backend.open_session(

tests/unit/test_session.py

@@ -62,18 +62,18 @@ def test_auth_args(self, mock_client_class):
 
         for args in connection_args:
             connection = databricks.sql.connect(**args)
-            host, port, http_path, *_ = mock_client_class.call_args[0]
-            self.assertEqual(args["server_hostname"], host)
-            self.assertEqual(args["http_path"], http_path)
+            call_kwargs = mock_client_class.call_args[1]
+            self.assertEqual(args["server_hostname"], call_kwargs["server_hostname"])
+            self.assertEqual(args["http_path"], call_kwargs["http_path"])
             connection.close()
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_http_header_passthrough(self, mock_client_class):
         http_headers = [("foo", "bar")]
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, http_headers=http_headers)
 
-        call_args = mock_client_class.call_args[0][3]
-        self.assertIn(("foo", "bar"), call_args)
+        call_kwargs = mock_client_class.call_args[1]
+        self.assertIn(("foo", "bar"), call_kwargs["http_headers"])
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_tls_arg_passthrough(self, mock_client_class):
@@ -95,7 +95,8 @@ def test_tls_arg_passthrough(self, mock_client_class):
     def test_useragent_header(self, mock_client_class):
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
 
-        http_headers = mock_client_class.call_args[0][3]
+        call_kwargs = mock_client_class.call_args[1]
+        http_headers = call_kwargs["http_headers"]
         user_agent_header = (
             "User-Agent",
             "{}/{}".format(databricks.sql.USER_AGENT_NAME, databricks.sql.__version__),
@@ -109,7 +110,8 @@ def test_useragent_header(self, mock_client_class):
                 databricks.sql.USER_AGENT_NAME, databricks.sql.__version__, "foobar"
             ),
         )
-        http_headers = mock_client_class.call_args[0][3]
+        call_kwargs = mock_client_class.call_args[1]
+        http_headers = call_kwargs["http_headers"]
         self.assertIn(user_agent_header_with_entry, http_headers)
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)