formatting: black + re-organise backend into new dir

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/db_client_interface.py renamed to src/databricks/sql/backend/databricks_client.py

@@ -43,11 +43,15 @@ def cancel_command(self, operation_handle: ttypes.TOperationHandle) -> None:
         pass
 
     @abstractmethod
-    def close_command(self, operation_handle: ttypes.TOperationHandle) -> ttypes.TStatus:
+    def close_command(
+        self, operation_handle: ttypes.TOperationHandle
+    ) -> ttypes.TStatus:
         pass
 
     @abstractmethod
-    def get_query_state(self, operation_handle: ttypes.TOperationHandle) -> ttypes.TOperationState:
+    def get_query_state(
+        self, operation_handle: ttypes.TOperationHandle
+    ) -> ttypes.TOperationState:
         pass
 
     @abstractmethod
@@ -132,4 +136,4 @@ def ssl_options(self) -> SSLOptions:
     @property
     @abstractmethod
     def max_download_threads(self) -> int:
-        pass 
+        pass

src/databricks/sql/thrift_backend.py renamed to src/databricks/sql/backend/thrift_backend.py

@@ -41,7 +41,7 @@
     convert_column_based_set_to_arrow_table,
 )
 from databricks.sql.types import SSLOptions
-from databricks.sql.db_client_interface import DatabricksClient
+from databricks.sql.backend.databricks_client import DatabricksClient
 
 logger = logging.getLogger(__name__)
 
@@ -221,11 +221,11 @@ def __init__(
     @property
     def staging_allowed_local_path(self) -> Union[None, str, List[str]]:
         return self._staging_allowed_local_path
-    
+
     @property
     def ssl_options(self) -> SSLOptions:
         return self._ssl_options
-        
+
     @property
     def max_download_threads(self) -> int:
         return self._max_download_threads
@@ -449,8 +449,10 @@ def attempt_request(attempt):
             except Exception as err:
                 error = err
                 retry_delay = extract_retry_delay(attempt)
-                error_message = ThriftDatabricksClient._extract_error_message_from_headers(
-                    getattr(self._transport, "headers", {})
+                error_message = (
+                    ThriftDatabricksClient._extract_error_message_from_headers(
+                        getattr(self._transport, "headers", {})
+                    )
                 )
             finally:
                 # Calling `close()` here releases the active HTTP connection back to the pool
@@ -710,7 +712,8 @@ def _col_to_description(col):
     @staticmethod
     def _hive_schema_to_description(t_table_schema):
         return [
-            ThriftDatabricksClient._col_to_description(col) for col in t_table_schema.columns
+            ThriftDatabricksClient._col_to_description(col)
+            for col in t_table_schema.columns
         ]
 
     def _results_message_to_execute_response(self, resp, operation_state):

src/databricks/sql/client.py

@@ -21,8 +21,8 @@
     CursorAlreadyClosedError,
 )
 from databricks.sql.thrift_api.TCLIService import ttypes
-from databricks.sql.thrift_backend import ThriftDatabricksClient
-from databricks.sql.db_client_interface import DatabricksClient
+from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
+from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.utils import (
     ExecuteResponse,
     ParamEscaper,
@@ -1230,7 +1230,7 @@ def _fill_results_buffer(self):
             )
 
         # Now we know self.backend is ThriftDatabricksClient, so it has fetch_results
-        thrift_backend_instance = self.backend # type: ThriftDatabricksClient
+        thrift_backend_instance = self.backend  # type: ThriftDatabricksClient
         results, has_more_rows = thrift_backend_instance.fetch_results(
             op_handle=self.command_id,
             max_rows=self.arraysize,
@@ -1447,7 +1447,8 @@ def close(self) -> None:
         If the connection has not been closed, and the cursor has not already
         been closed on the server for some other reason, issue a request to the server to close it.
         """
-        # TODO: the state is still thrift specific, define some ENUM for status that each service has to map to 
+        # TODO: the state is still thrift specific, define some ENUM for status that each service has to map to
+        # when we generalise the ResultSet
         try:
             if (
                 self.op_state != ttypes.TOperationState.CLOSED_STATE

src/databricks/sql/session.py

@@ -7,8 +7,8 @@
 from databricks.sql.exc import SessionAlreadyClosedError, DatabaseError, RequestError
 from databricks.sql import __version__
 from databricks.sql import USER_AGENT_NAME
-from databricks.sql.thrift_backend import ThriftDatabricksClient
-from databricks.sql.db_client_interface import DatabricksClient
+from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
+from databricks.sql.backend.databricks_client import DatabricksClient
 
 logger = logging.getLogger(__name__)
 

tests/unit/test_client.py

@@ -15,7 +15,7 @@
     THandleIdentifier,
     TOperationType,
 )
-from databricks.sql.thrift_backend import ThriftDatabricksClient
+from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 
 import databricks.sql
 import databricks.sql.client as client

tests/unit/test_fetches.py

@@ -9,7 +9,7 @@
 
 import databricks.sql.client as client
 from databricks.sql.utils import ExecuteResponse, ArrowQueue
-from databricks.sql.thrift_backend import ThriftDatabricksClient
+from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 
 
 @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")

tests/unit/test_thrift_backend.py

@@ -17,7 +17,7 @@
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql import *
 from databricks.sql.auth.authenticators import AuthProvider
-from databricks.sql.thrift_backend import ThriftDatabricksClient
+from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 
 
 def retry_policy_factory():