fix: databricks import lint

dlt-hub · Dec 5, 2024 · f75b8c8 · f75b8c8
1 parent 9c25d72
commit f75b8c8
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/dlt/destinations/impl/databricks/sql_client.py b/dlt/destinations/impl/databricks/sql_client.py
@@ -13,7 +13,8 @@
     Dict,
 )
 
-import databricks.sql
+from databricks import sql as databricks_lib  # type: ignore[attr-defined]
+
 from databricks.sql.client import (
     Connection as DatabricksSqlConnection,
     Cursor as DatabricksSqlCursor,
@@ -59,7 +60,7 @@ def iter_df(self, chunk_size: int) -> Generator[DataFrame, None, None]:
 
 
 class DatabricksSqlClient(SqlClientBase[DatabricksSqlConnection], DBTransaction):
-    dbapi: ClassVar[DBApi] = databricks.sql
+    dbapi: ClassVar[DBApi] = databricks_lib
 
     def __init__(
         self,
@@ -74,7 +75,7 @@ def __init__(
 
     def open_connection(self) -> DatabricksSqlConnection:
         conn_params = self.credentials.to_connector_params()
-        self._conn = databricks.sql.connect(
+        self._conn = databricks_lib.sql.connect(
             **conn_params, schema=self.dataset_name, use_inline_params="silent"
         )
         return self._conn
@@ -155,23 +156,23 @@ def catalog_name(self, escape: bool = True) -> Optional[str]:
 
     @staticmethod
     def _make_database_exception(ex: Exception) -> Exception:
-        if isinstance(ex, databricks.sql.ServerOperationError):
+        if isinstance(ex, databricks_lib.ServerOperationError):
             if "TABLE_OR_VIEW_NOT_FOUND" in str(ex):
                 return DatabaseUndefinedRelation(ex)
             elif "SCHEMA_NOT_FOUND" in str(ex):
                 return DatabaseUndefinedRelation(ex)
             elif "PARSE_SYNTAX_ERROR" in str(ex):
                 return DatabaseTransientException(ex)
             return DatabaseTerminalException(ex)
-        elif isinstance(ex, databricks.sql.OperationalError):
+        elif isinstance(ex, databricks_lib.OperationalError):
             return DatabaseTransientException(ex)
-        elif isinstance(ex, (databricks.sql.ProgrammingError, databricks.sql.IntegrityError)):
+        elif isinstance(ex, (databricks_lib.ProgrammingError, databricks_lib.IntegrityError)):
             return DatabaseTerminalException(ex)
-        elif isinstance(ex, databricks.sql.DatabaseError):
+        elif isinstance(ex, databricks_lib.DatabaseError):
             return DatabaseTransientException(ex)
         else:
             return DatabaseTransientException(ex)
 
     @staticmethod
     def is_dbapi_exception(ex: Exception) -> bool:
-        return isinstance(ex, databricks.sql.DatabaseError)
+        return isinstance(ex, databricks_lib.DatabaseError)