fix the integration workflow and linting

Author: benc-db

.github/workflows/integration.yml

@@ -62,4 +62,4 @@ jobs:
       - name: Run e2e tests
         run: poetry run python -m pytest tests/e2e
       - name: Run SQL Alchemy tests
-        run: poetry run python -m pytest sqlalchemy/tests/test_local
+        run: poetry run python -m pytest src/databricks/sqlalchemy/tests/test_local

src/databricks/sql/client.py

@@ -205,7 +205,9 @@ def read(self) -> Optional[OAuthToken]:
         self.disable_pandas = kwargs.get("_disable_pandas", False)
         self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True)
 
-        auth_provider = get_python_sql_connector_auth_provider(server_hostname, **kwargs)
+        auth_provider = get_python_sql_connector_auth_provider(
+            server_hostname, **kwargs
+        )
 
         if not kwargs.get("_user_agent_entry"):
             useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__)
@@ -276,7 +278,8 @@ def __exit__(self, exc_type, exc_value, traceback):
     def __del__(self):
         if self.open:
             logger.debug(
-                "Closing unclosed connection for session " "{}".format(self.get_session_id_hex())
+                "Closing unclosed connection for session "
+                "{}".format(self.get_session_id_hex())
             )
             try:
                 self._close(close_cursors=False)
@@ -356,9 +359,13 @@ def _close(self, close_cursors=True) -> None:
                 logger.info("Session was closed by a prior request")
         except DatabaseError as e:
             if "Invalid SessionHandle" in str(e):
-                logger.warning(f"Attempted to close session that was already closed: {e}")
+                logger.warning(
+                    f"Attempted to close session that was already closed: {e}"
+                )
             else:
-                logger.warning(f"Attempt to close session raised an exception at the server: {e}")
+                logger.warning(
+                    f"Attempt to close session raised an exception at the server: {e}"
+                )
         except Exception as e:
             logger.error(f"Attempt to close session raised a local exception: {e}")
 
@@ -441,7 +448,9 @@ def _all_dbsql_parameters_are_named(self, params: List[TDbsqlParameter]) -> bool
         """Return True if all members of the list have a non-null .name attribute"""
         return all([i.name is not None for i in params])
 
-    def _normalize_tparametersequence(self, params: TParameterSequence) -> List[TDbsqlParameter]:
+    def _normalize_tparametersequence(
+        self, params: TParameterSequence
+    ) -> List[TDbsqlParameter]:
         """Retains the same order as the input list."""
 
         output: List[TDbsqlParameter] = []
@@ -453,9 +462,12 @@ def _normalize_tparametersequence(self, params: TParameterSequence) -> List[TDbs
 
         return output
 
-    def _normalize_tparameterdict(self, params: TParameterDict) -> List[TDbsqlParameter]:
+    def _normalize_tparameterdict(
+        self, params: TParameterDict
+    ) -> List[TDbsqlParameter]:
         return [
-            dbsql_parameter_from_primitive(value=value, name=name) for name, value in params.items()
+            dbsql_parameter_from_primitive(value=value, name=name)
+            for name, value in params.items()
         ]
 
     def _normalize_tparametercollection(
@@ -528,7 +540,8 @@ def _prepare_native_parameters(
 
         stmt = stmt
         output = [
-            p.as_tspark_param(named=param_structure == ParameterStructure.NAMED) for p in params
+            p.as_tspark_param(named=param_structure == ParameterStructure.NAMED)
+            for p in params
         ]
 
         return stmt, output
@@ -544,7 +557,9 @@ def _check_not_closed(self):
         if not self.open:
             raise Error("Attempting operation on closed cursor")
 
-    def _handle_staging_operation(self, staging_allowed_local_path: Union[None, str, List[str]]):
+    def _handle_staging_operation(
+        self, staging_allowed_local_path: Union[None, str, List[str]]
+    ):
         """Fetch the HTTP request instruction from a staging ingestion command
         and call the designated handler.
 
@@ -561,7 +576,9 @@ def _handle_staging_operation(self, staging_allowed_local_path: Union[None, str,
                 "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands"
             )
 
-        abs_staging_allowed_local_paths = [os.path.abspath(i) for i in _staging_allowed_local_paths]
+        abs_staging_allowed_local_paths = [
+            os.path.abspath(i) for i in _staging_allowed_local_paths
+        ]
 
         assert self.active_result_set is not None
         row = self.active_result_set.fetchone()
@@ -589,7 +606,9 @@ def _handle_staging_operation(self, staging_allowed_local_path: Union[None, str,
                 )
 
         # May be real headers, or could be json string
-        headers = json.loads(row.headers) if isinstance(row.headers, str) else row.headers
+        headers = (
+            json.loads(row.headers) if isinstance(row.headers, str) else row.headers
+        )
 
         handler_args = {
             "presigned_url": row.presignedUrl,
@@ -616,7 +635,9 @@ def _handle_staging_operation(self, staging_allowed_local_path: Union[None, str,
                 + "Supported operations are GET, PUT, and REMOVE"
             )
 
-    def _handle_staging_put(self, presigned_url: str, local_file: str, headers: dict = None):
+    def _handle_staging_put(
+        self, presigned_url: str, local_file: str, headers: dict = None
+    ):
         """Make an HTTP PUT request
 
         Raise an exception if request fails. Returns no data.
@@ -639,15 +660,19 @@ def _handle_staging_put(self, presigned_url: str, local_file: str, headers: dict
         # fmt: on
 
         if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
-            raise Error(f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}")
+            raise Error(
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+            )
 
         if r.status_code == ACCEPTED:
             logger.debug(
                 f"Response code {ACCEPTED} from server indicates ingestion command was accepted "
                 + "but not yet applied on the server. It's possible this command may fail later."
             )
 
-    def _handle_staging_get(self, local_file: str, presigned_url: str, headers: dict = None):
+    def _handle_staging_get(
+        self, local_file: str, presigned_url: str, headers: dict = None
+    ):
         """Make an HTTP GET request, create a local file with the received data
 
         Raise an exception if request fails. Returns no data.
@@ -661,7 +686,9 @@ def _handle_staging_get(self, local_file: str, presigned_url: str, headers: dict
         # response.ok verifies the status code is not between 400-600.
         # Any 2xx or 3xx will evaluate r.ok == True
         if not r.ok:
-            raise Error(f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}")
+            raise Error(
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+            )
 
         with open(local_file, "wb") as fp:
             fp.write(r.content)
@@ -672,7 +699,9 @@ def _handle_staging_remove(self, presigned_url: str, headers: dict = None):
         r = requests.delete(url=presigned_url, headers=headers)
 
         if not r.ok:
-            raise Error(f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}")
+            raise Error(
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+            )
 
     def execute(
         self,
@@ -967,7 +996,8 @@ def cancel(self) -> None:
             self.thrift_backend.cancel_command(self.active_op_handle)
         else:
             logger.warning(
-                "Attempting to cancel a command, but there is no " "currently executing command"
+                "Attempting to cancel a command, but there is no "
+                "currently executing command"
             )
 
     def close(self) -> None:
@@ -1085,7 +1115,9 @@ def _convert_arrow_table(self, table):
         ResultRow = Row(*column_names)
 
         if self.connection.disable_pandas is True:
-            return [ResultRow(*[v.as_py() for v in r]) for r in zip(*table.itercolumns())]
+            return [
+                ResultRow(*[v.as_py() for v in r]) for r in zip(*table.itercolumns())
+            ]
 
         # Need to use nullable types, as otherwise type can change when there are missing values.
         # See https://arrow.apache.org/docs/python/pandas.html#nullable-types
@@ -1132,7 +1164,11 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table:
         n_remaining_rows = size - results.num_rows
         self._next_row_index += results.num_rows
 
-        while n_remaining_rows > 0 and not self.has_been_closed_server_side and self.has_more_rows:
+        while (
+            n_remaining_rows > 0
+            and not self.has_been_closed_server_side
+            and self.has_more_rows
+        ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
             results = pyarrow.concat_tables([results, partial_results])