lint: fixed static type check issue

datahub-project · Dec 26, 2024 · 07f565f · 07f565f
1 parent e2bc694
commit 07f565f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/source.py b/metadata-ingestion/src/datahub/ingestion/source/sql/mssql/source.py
@@ -724,7 +724,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
             ):
                 yield from auto_workunit(
                     generate_procedure_lineage(
-                        schema_resolver=self.schema_resolver,
+                        schema_resolver=self.get_schema_resolver(),
                         procedure=procedure,
                         procedure_job_urn=MSSQLDataJob(entity=procedure).urn,
                         is_temp_table=self.is_temp_table,

diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
@@ -105,6 +105,7 @@
     TagAssociationClass,
     ViewPropertiesClass,
 )
+from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.sql_parsing.sql_parsing_aggregator import SqlParsingAggregator
 from datahub.telemetry import telemetry
 from datahub.utilities.registries.domain_registry import DomainRegistry
@@ -1349,5 +1350,8 @@ def prepare_profiler_args(
             schema=schema, table=table, partition=partition, custom_sql=custom_sql
         )
 
+    def get_schema_resolver(self) -> SchemaResolver:
+        return self.aggregator._schema_resolver
+
     def get_report(self):
         return self.report