MINOR: Fix databricks connector (#15894)

This commit is contained in:
Mayur Singal 2024-04-15 16:18:58 +05:30 committed by GitHub
parent c0a58c5b44
commit 269ed6d5bb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -20,7 +20,6 @@ from sqlalchemy import types, util
from sqlalchemy.engine import reflection
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.exc import DatabaseError
from sqlalchemy.inspection import inspect
from sqlalchemy.sql.sqltypes import String
from sqlalchemy_databricks._dialect import DatabricksDialect
@ -310,8 +309,9 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
new_service_connection = deepcopy(self.service_connection)
new_service_connection.catalog = database_name
self.engine = get_connection(new_service_connection)
self.inspector = inspect(self.engine)
self._connection = None # Lazy init as well
self._connection_map = {} # Lazy init as well
self._inspector_map = {}
def get_configured_database(self) -> Optional[str]:
return self.service_connection.catalog
@ -598,7 +598,7 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
DATABRICKS_GET_TABLE_COMMENTS.format(
schema_name=schema_name,
table_name=table_name,
catalog_name=self.context.database,
catalog_name=self.context.get().database,
)
)
for result in list(cursor):
@ -607,7 +607,7 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
description = data[1] if data and data[1] else None
elif data[0] and data[0].strip() == "Location":
self.external_location_map[
(self.context.database, schema_name, table_name)
(self.context.get().database, schema_name, table_name)
] = (
data[1]
if data and data[1] and not data[1].startswith("dbfs")