Resolve circular dependencies

Signed-off-by: Levko Kravets <[email protected]>

Author: kravets-levko

src/databricks/sql/auth/thrift_http_client.py

@@ -14,7 +14,7 @@
 from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager
 from urllib3.util import make_headers
 from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy
-from databricks.sql.utils import SSLOptions
+from databricks.sql.types import SSLOptions
 
 logger = logging.getLogger(__name__)
 

src/databricks/sql/client.py

@@ -22,7 +22,6 @@
     ParamEscaper,
     inject_parameters,
     transform_paramstyle,
-    SSLOptions,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -36,7 +35,7 @@
 )
 
 
-from databricks.sql.types import Row
+from databricks.sql.types import Row, SSLOptions
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
 from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 

src/databricks/sql/cloudfetch/download_manager.py

@@ -1,6 +1,5 @@
 import logging
 
-from ssl import SSLContext
 from concurrent.futures import ThreadPoolExecutor, Future
 from typing import List, Union
 
@@ -9,7 +8,7 @@
     DownloadableResultSettings,
     DownloadedFile,
 )
-from databricks.sql.utils import SSLOptions
+from databricks.sql.types import SSLOptions
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 

src/databricks/sql/cloudfetch/downloader.py

@@ -8,7 +8,7 @@
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.exc import Error
-from databricks.sql.utils import SSLOptions
+from databricks.sql.types import SSLOptions
 
 logger = logging.getLogger(__name__)
 

src/databricks/sql/thrift_backend.py

@@ -5,7 +5,6 @@
 import time
 import uuid
 import threading
-from ssl import CERT_NONE, CERT_REQUIRED, create_default_context
 from typing import List, Union
 
 import pyarrow
@@ -35,8 +34,8 @@
     convert_arrow_based_set_to_arrow_table,
     convert_decimals_in_arrow_table,
     convert_column_based_set_to_arrow_table,
-    SSLOptions,
 )
+from databricks.sql.types import SSLOptions
 
 logger = logging.getLogger(__name__)
 

src/databricks/sql/types.py

@@ -19,6 +19,54 @@
 from typing import Any, Dict, List, Optional, Tuple, Union, TypeVar
 import datetime
 import decimal
+from ssl import SSLContext, CERT_NONE, CERT_REQUIRED, create_default_context
+
+
+class SSLOptions:
+    tls_verify: bool
+    tls_verify_hostname: bool
+    tls_trusted_ca_file: Optional[str]
+    tls_client_cert_file: Optional[str]
+    tls_client_cert_key_file: Optional[str]
+    tls_client_cert_key_password: Optional[str]
+
+    def __init__(
+        self,
+        tls_verify: Optional[bool] = True,
+        tls_verify_hostname: Optional[bool] = True,
+        tls_trusted_ca_file: Optional[str] = None,
+        tls_client_cert_file: Optional[str] = None,
+        tls_client_cert_key_file: Optional[str] = None,
+        tls_client_cert_key_password: Optional[str] = None,
+    ):
+        self.tls_verify = tls_verify
+        self.tls_verify_hostname = tls_verify_hostname
+        self.tls_trusted_ca_file = tls_trusted_ca_file
+        self.tls_client_cert_file = tls_client_cert_file
+        self.tls_client_cert_key_file = tls_client_cert_key_file
+        self.tls_client_cert_key_password = tls_client_cert_key_password
+
+    def create_ssl_context(self) -> SSLContext:
+        ssl_context = create_default_context(cafile=self.tls_trusted_ca_file)
+
+        if self.tls_verify is False:
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = CERT_NONE
+        elif self.tls_verify_hostname is False:
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = CERT_REQUIRED
+        else:
+            ssl_context.check_hostname = True
+            ssl_context.verify_mode = CERT_REQUIRED
+
+        if self.tls_client_cert_file:
+            ssl_context.load_cert_chain(
+                certfile=self.tls_client_cert_file,
+                keyfile=self.tls_client_cert_key_file,
+                password=self.tls_client_cert_key_password,
+            )
+
+        return ssl_context
 
 
 class Row(tuple):

src/databricks/sql/utils.py

@@ -9,7 +9,6 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 import re
-from ssl import SSLContext, CERT_NONE, CERT_REQUIRED, create_default_context
 
 import lz4.frame
 import pyarrow
@@ -21,6 +20,7 @@
     TSparkArrowResultLink,
     TSparkRowSetType,
 )
+from databricks.sql.types import SSLOptions
 
 from databricks.sql.parameters.native import ParameterStructure, TDbsqlParameter
 
@@ -31,53 +31,6 @@
 logger = logging.getLogger(__name__)
 
 
-class SSLOptions:
-    tls_verify: bool
-    tls_verify_hostname: bool
-    tls_trusted_ca_file: Optional[str]
-    tls_client_cert_file: Optional[str]
-    tls_client_cert_key_file: Optional[str]
-    tls_client_cert_key_password: Optional[str]
-
-    def __init__(
-        self,
-        tls_verify: Optional[bool] = True,
-        tls_verify_hostname: Optional[bool] = True,
-        tls_trusted_ca_file: Optional[str] = None,
-        tls_client_cert_file: Optional[str] = None,
-        tls_client_cert_key_file: Optional[str] = None,
-        tls_client_cert_key_password: Optional[str] = None,
-    ):
-        self.tls_verify = tls_verify
-        self.tls_verify_hostname = tls_verify_hostname
-        self.tls_trusted_ca_file = tls_trusted_ca_file
-        self.tls_client_cert_file = tls_client_cert_file
-        self.tls_client_cert_key_file = tls_client_cert_key_file
-        self.tls_client_cert_key_password = tls_client_cert_key_password
-
-    def create_ssl_context(self) -> SSLContext:
-        ssl_context = create_default_context(cafile=self.tls_trusted_ca_file)
-
-        if self.tls_verify is False:
-            ssl_context.check_hostname = False
-            ssl_context.verify_mode = CERT_NONE
-        elif self.tls_verify_hostname is False:
-            ssl_context.check_hostname = False
-            ssl_context.verify_mode = CERT_REQUIRED
-        else:
-            ssl_context.check_hostname = True
-            ssl_context.verify_mode = CERT_REQUIRED
-
-        if self.tls_client_cert_file:
-            ssl_context.load_cert_chain(
-                certfile=self.tls_client_cert_file,
-                keyfile=self.tls_client_cert_key_file,
-                password=self.tls_client_cert_key_password,
-            )
-
-        return ssl_context
-
-
 class ResultSetQueue(ABC):
     @abstractmethod
     def next_n_rows(self, num_rows: int) -> pyarrow.Table: