openmetadata-ingestion 1.3.3.0rc2__py3-none-any.whl → 1.4.0.0rc1__py3-none-any.whl

metadata/ingestion/source/database/database_service.py

@@ -62,11 +62,12 @@ from metadata.ingestion.models.ometa_classification import OMetaTagAndClassifica
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.source.connections import get_test_connection_fn
 from metadata.utils import fqn
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.filters import filter_by_schema
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.tag_utils import get_tag_label
@@ -108,7 +109,11 @@ class DatabaseServiceTopology(ServiceTopology):
         # Note how we have `yield_view_lineage` and `yield_stored_procedure_lineage`
         # as post_processed. This is because we cannot ensure proper lineage processing
         # until we have finished ingesting all the metadata from the source.
-        post_process=["yield_view_lineage", "yield_procedure_lineage_and_queries"],
+        post_process=[
+            "yield_view_lineage",
+            "yield_procedure_lineage_and_queries",
+            "yield_external_table_lineage",
+        ],
     )
     database = TopologyNode(
         producer="get_database_names",
@@ -152,6 +157,7 @@ class DatabaseServiceTopology(ServiceTopology):
         ],
         children=["table", "stored_procedure"],
         post_process=["mark_tables_as_deleted", "mark_stored_procedures_as_deleted"],
+        threads=True,
     )
     table = TopologyNode(
         producer="get_tables_name_and_type",
@@ -212,7 +218,7 @@ class DatabaseServiceSource(
     inspector: Inspector
 
     topology = DatabaseServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
 
     @property
     def name(self) -> str:
@@ -380,7 +386,7 @@ class DatabaseServiceSource(
         """
 
         tag_labels = []
-        for tag_and_category in self.context.tags or []:
+        for tag_and_category in self.context.get().tags or []:
             if tag_and_category.fqn and tag_and_category.fqn.__root__ == entity_fqn:
                 tag_label = get_tag_label(
                     metadata=self.metadata,
@@ -400,7 +406,7 @@ class DatabaseServiceSource(
         database_fqn = fqn.build(
             self.metadata,
             entity_type=Database,
-            service_name=self.context.database_service,
+            service_name=self.context.get().database_service,
             database_name=database_name,
         )
         return self.get_tag_by_fqn(entity_fqn=database_fqn)
@@ -414,12 +420,13 @@ class DatabaseServiceSource(
         schema_fqn = fqn.build(
             self.metadata,
             entity_type=DatabaseSchema,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
+            service_name=self.context.get().database_service,
+            database_name=self.context.get().database,
             schema_name=schema_name,
         )
         return self.get_tag_by_fqn(entity_fqn=schema_fqn)
 
+    @calculate_execution_time()
     def get_tag_labels(self, table_name: str) -> Optional[List[TagLabel]]:
         """
         This will only get executed if the tags context
@@ -428,9 +435,9 @@ class DatabaseServiceSource(
         table_fqn = fqn.build(
             self.metadata,
             entity_type=Table,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
+            service_name=self.context.get().database_service,
+            database_name=self.context.get().database,
+            schema_name=self.context.get().database_schema,
             table_name=table_name,
             skip_es_search=True,
         )
@@ -446,14 +453,15 @@ class DatabaseServiceSource(
         col_fqn = fqn.build(
             self.metadata,
             entity_type=Column,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
+            service_name=self.context.get().database_service,
+            database_name=self.context.get().database,
+            schema_name=self.context.get().database_schema,
             table_name=table_name,
             column_name=column["name"],
         )
         return self.get_tag_by_fqn(entity_fqn=col_fqn)
 
+    @calculate_execution_time()
     def register_record(self, table_request: CreateTableRequest) -> None:
         """
         Mark the table record as scanned and update the database_source_state
@@ -461,9 +469,9 @@ class DatabaseServiceSource(
         table_fqn = fqn.build(
             self.metadata,
             entity_type=Table,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
+            service_name=self.context.get().database_service,
+            database_name=self.context.get().database,
+            schema_name=self.context.get().database_schema,
             table_name=table_request.name.__root__,
             skip_es_search=True,
         )
@@ -479,9 +487,9 @@ class DatabaseServiceSource(
         table_fqn = fqn.build(
             self.metadata,
             entity_type=StoredProcedure,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
+            service_name=self.context.get().database_service,
+            database_name=self.context.get().database,
+            schema_name=self.context.get().database_schema,
             procedure_name=stored_proc_request.name.__root__,
         )
 
@@ -494,8 +502,8 @@ class DatabaseServiceSource(
             schema_fqn = fqn.build(
                 self.metadata,
                 entity_type=DatabaseSchema,
-                service_name=self.context.database_service,
-                database_name=self.context.database,
+                service_name=self.context.get().database_service,
+                database_name=self.context.get().database,
                 schema_name=schema_name,
             )
             if filter_by_schema(
@@ -507,6 +515,7 @@ class DatabaseServiceSource(
                 continue
             yield schema_fqn if return_fqn else schema_name
 
+    @calculate_execution_time()
     def get_owner_ref(self, table_name: str) -> Optional[EntityReference]:
         """
         Method to process the table owners
@@ -514,9 +523,9 @@ class DatabaseServiceSource(
         try:
             if self.source_config.includeOwners:
                 owner_name = self.inspector.get_table_owner(
-                    connection=self.connection,  # pylint: disable=no-member
+                    connection=self.connection,  # pylint: disable=no-member.fetchall()
                     table_name=table_name,
-                    schema=self.context.database_schema,
+                    schema=self.context.get().database_schema,
                 )
                 owner_ref = self.metadata.get_reference_by_name(name=owner_name)
                 return owner_ref
@@ -529,14 +538,14 @@ class DatabaseServiceSource(
         """
         Use the current inspector to mark tables as deleted
         """
-        if not self.context.__dict__.get("database"):
+        if not self.context.get().__dict__.get("database"):
             raise ValueError(
                 "No Database found in the context. We cannot run the table deletion."
             )
 
         if self.source_config.markDeletedTables:
             logger.info(
-                f"Mark Deleted Tables set to True. Processing database [{self.context.database}]"
+                f"Mark Deleted Tables set to True. Processing database [{self.context.get().database}]"
             )
             schema_fqn_list = self._get_filtered_schema_names(
                 return_fqn=True, add_to_status=False
@@ -557,7 +566,7 @@ class DatabaseServiceSource(
         """
         if self.source_config.markDeletedStoredProcedures:
             logger.info(
-                f"Mark Deleted Stored Procedures Processing database [{self.context.database}]"
+                f"Mark Deleted Stored Procedures Processing database [{self.context.get().database}]"
             )
 
             schema_fqn_list = self._get_filtered_schema_names(
@@ -578,6 +587,11 @@ class DatabaseServiceSource(
         Get the life cycle data of the table
         """
 
+    def yield_external_table_lineage(self) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Process external table lineage
+        """
+
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
         test_connection_fn(self.metadata, self.connection_obj, self.service_connection)

metadata/ingestion/source/database/databricks/metadata.py

@@ -18,8 +18,8 @@ from typing import Iterable, Optional, Tuple, Union
 from pyhive.sqlalchemy_hive import _type_map
 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
 
@@ -51,6 +51,9 @@ from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_TABLE_TAGS,
     DATABRICKS_VIEW_DEFINITIONS,
 )
+from metadata.ingestion.source.database.external_table_lineage_mixin import (
+    ExternalTableLineageMixin,
+)
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.utils import fqn
 from metadata.utils.constants import DEFAULT_DATABASE
@@ -258,7 +261,7 @@ DatabricksDialect.get_all_view_definitions = get_all_view_definitions
 reflection.Inspector.get_schema_names = get_schema_names_reflection
 
 
-class DatabricksSource(CommonDbSourceService, MultiDBSource):
+class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from Databricks Source using
@@ -272,6 +275,7 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         self.catalog_tags = {}
         self.schema_tags = {}
         self.table_tags = {}
+        self.external_location_map = {}
         self.column_tags = {}
 
     def _init_version(self):
@@ -283,7 +287,9 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
             self.is_older_version = True
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: DatabricksConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, DatabricksConnection):
@@ -303,8 +309,9 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         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
@@ -422,7 +429,7 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                 database_fqn = fqn.build(
                     self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=new_catalog,
                 )
                 if filter_by_database(
@@ -448,7 +455,7 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
             yield self.service_connection.databaseSchema
         else:
             for schema_name in self.inspector.get_schema_names(
-                database=self.context.database,
+                database=self.context.get().database,
                 is_old_version=self.is_older_version,
             ):
                 yield schema_name
@@ -466,7 +473,7 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                     tag_fqn=fqn.build(
                         self.metadata,
                         Database,
-                        service_name=self.context.database_service,
+                        service_name=self.context.get().database_service,
                         database_name=database_name,
                     ),
                     tags=[tag_value],
@@ -491,14 +498,16 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         Method to yield schema tags
         """
         try:
-            schema_tags = self.schema_tags.get((self.context.database, schema_name), [])
+            schema_tags = self.schema_tags.get(
+                (self.context.get().database, schema_name), []
+            )
             for tag_name, tag_value in schema_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
                         DatabaseSchema,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
                         schema_name=schema_name,
                     ),
                     tags=[tag_value],
@@ -522,16 +531,21 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         table_name, _ = table_name_and_type
         try:
             table_tags = self.table_tags.get(
-                (self.context.database, self.context.database_schema, table_name), []
+                (
+                    self.context.get().database,
+                    self.context.get().database_schema,
+                    table_name,
+                ),
+                [],
             )
             for tag_name, tag_value in table_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
                         Table,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
+                        schema_name=self.context.get().database_schema,
                         table_name=table_name,
                     ),
                     tags=[tag_value],
@@ -541,7 +555,12 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                 )
 
             column_tags = self.column_tags.get(
-                (self.context.database, self.context.database_schema, table_name), {}
+                (
+                    self.context.get().database,
+                    self.context.get().database_schema,
+                    table_name,
+                ),
+                {},
             )
             for column_name, tags in column_tags.items():
                 for tag_name, tag_value in tags or []:
@@ -549,9 +568,9 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         tag_fqn=fqn.build(
                             self.metadata,
                             Column,
-                            service_name=self.context.database_service,
-                            database_name=self.context.database,
-                            schema_name=self.context.database_schema,
+                            service_name=self.context.get().database_service,
+                            database_name=self.context.get().database,
+                            schema_name=self.context.get().database_schema,
                             table_name=table_name,
                             column_name=column_name,
                         ),
@@ -569,3 +588,36 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                     stackTrace=traceback.format_exc(),
                 )
             )
+
+    def get_table_description(
+        self, schema_name: str, table_name: str, inspector: Inspector
+    ) -> str:
+        description = None
+        try:
+            cursor = self.connection.execute(
+                DATABRICKS_GET_TABLE_COMMENTS.format(
+                    schema_name=schema_name,
+                    table_name=table_name,
+                    catalog_name=self.context.get().database,
+                )
+            )
+            for result in list(cursor):
+                data = result.values()
+                if data[0] and data[0].strip() == "Comment":
+                    description = data[1] if data and data[1] else None
+                elif data[0] and data[0].strip() == "Location":
+                    self.external_location_map[
+                        (self.context.get().database, schema_name, table_name)
+                    ] = (
+                        data[1]
+                        if data and data[1] and not data[1].startswith("dbfs")
+                        else None
+                    )
+
+        # Catch any exception without breaking the ingestion
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Table description error for table [{schema_name}.{table_name}]: {exc}"
+            )
+        return description

metadata/ingestion/source/database/databricks/queries.py

@@ -25,7 +25,9 @@ DATABRICKS_VIEW_DEFINITIONS = textwrap.dedent(
 )
 
 
-DATABRICKS_GET_TABLE_COMMENTS = "DESCRIBE TABLE EXTENDED {schema_name}.{table_name}"
+DATABRICKS_GET_TABLE_COMMENTS = (
+    "DESCRIBE TABLE EXTENDED {catalog_name}.{schema_name}.{table_name}"
+)
 
 DATABRICKS_GET_CATALOGS = "SHOW CATALOGS"
 

metadata/ingestion/source/database/databricks/query_parser.py

@@ -12,6 +12,7 @@
 Databricks Query parser module
 """
 from abc import ABC
+from typing import Optional
 
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
@@ -35,7 +36,11 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
 
     filters: str
 
-    def _init_super(self, config: WorkflowSource, metadata: OpenMetadata):
+    def _init_super(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
         super().__init__(config, metadata, False)
 
     # pylint: disable=super-init-not-called
@@ -44,7 +49,9 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
         self.client = DatabricksClient(self.service_connection)
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: DatabricksConnection = config.serviceConnection.__root__.config

metadata/ingestion/source/database/datalake/connection.py

@@ -89,7 +89,6 @@ def _(config: GCSConfig):
 
 @get_datalake_client.register
 def _(config: AzureConfig):
-
     try:
         return AzureClient(config.securityConfig).create_blob_client()
     except Exception as exc:
@@ -104,7 +103,7 @@ def set_gcs_datalake_client(config: GCSConfig, project_id: str):
         gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
             project_id
         )
-    return get_datalake_client(config=gcs_config)
+    return get_datalake_client(gcs_config)
 
 
 def get_connection(connection: DatalakeConnection) -> DatalakeClient:

metadata/ingestion/source/database/datalake/metadata.py

@@ -15,7 +15,7 @@ DataLake connector to fetch metadata from a files stored s3, gcs and Hdfs
 import json
 import os
 import traceback
-from typing import Any, Iterable, Tuple, Union
+from typing import Any, Iterable, Optional, Tuple, Union
 
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
@@ -117,7 +117,9 @@ class DatalakeSource(DatabaseServiceSource):
         self.reader = get_reader(config_source=self.config_source, client=self.client)
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: DatalakeConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, DatalakeConnection):
@@ -148,7 +150,7 @@ class DatalakeSource(DatabaseServiceSource):
                 database_fqn = fqn.build(
                     self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=project_id,
                 )
                 if filter_by_database(
@@ -189,7 +191,7 @@ class DatalakeSource(DatabaseServiceSource):
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -204,8 +206,8 @@ class DatalakeSource(DatabaseServiceSource):
                 schema_fqn = fqn.build(
                     self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                     schema_name=bucket.name,
                 )
 
@@ -236,8 +238,8 @@ class DatalakeSource(DatabaseServiceSource):
             schema_fqn = fqn.build(
                 self.metadata,
                 entity_type=DatabaseSchema,
-                service_name=self.context.database_service,
-                database_name=self.context.database,
+                service_name=self.context.get().database_service,
+                database_name=self.context.get().database,
                 schema_name=bucket["Name"],
             )
             if filter_by_schema(
@@ -284,8 +286,8 @@ class DatalakeSource(DatabaseServiceSource):
             schema_fqn = fqn.build(
                 self.metadata,
                 entity_type=DatabaseSchema,
-                service_name=self.context.database_service,
-                database_name=self.context.database,
+                service_name=self.context.get().database_service,
+                database_name=self.context.get().database,
                 schema_name=schema["name"],
             )
             if filter_by_schema(
@@ -312,8 +314,8 @@ class DatalakeSource(DatabaseServiceSource):
                 database=fqn.build(
                     metadata=self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                 ),
             )
         )
@@ -329,7 +331,7 @@ class DatalakeSource(DatabaseServiceSource):
 
         :return: tables or views, depending on config
         """
-        bucket_name = self.context.database_schema
+        bucket_name = self.context.get().database_schema
         prefix = self.service_connection.prefix
         try:
             metadata_config_response = self.reader.read(
@@ -404,7 +406,7 @@ class DatalakeSource(DatabaseServiceSource):
         Prepare a table request and pass it to the sink
         """
         table_name, table_type, table_extension = table_name_and_type
-        schema_name = self.context.database_schema
+        schema_name = self.context.get().database_schema
         try:
             table_constraints = None
             data_frame, raw_data = fetch_dataframe(
@@ -434,8 +436,8 @@ class DatalakeSource(DatabaseServiceSource):
                     databaseSchema=fqn.build(
                         metadata=self.metadata,
                         entity_type=DatabaseSchema,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
                         schema_name=schema_name,
                     ),
                     fileFormat=table_extension.value if table_extension else None,
@@ -486,9 +488,9 @@ class DatalakeSource(DatabaseServiceSource):
         table_fqn = fqn.build(
             self.metadata,
             entity_type=Table,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
+            service_name=self.context.get().database_service,
+            database_name=self.context.get().database,
+            schema_name=self.context.get().database_schema,
             table_name=table_name,
             skip_es_search=True,
         )

metadata/ingestion/source/database/db2/metadata.py

@@ -10,6 +10,7 @@
 #  limitations under the License.
 """Db2 source module"""
 import traceback
+from typing import Optional
 
 from ibm_db_sa.base import DB2Dialect
 from sqlalchemy.engine import reflection
@@ -47,7 +48,9 @@ class Db2Source(CommonDbSourceService):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: Db2Connection = config.serviceConnection.__root__.config
         if not isinstance(connection, Db2Connection):

metadata/ingestion/source/database/dbt/dbt_service.py

@@ -31,7 +31,7 @@ from metadata.ingestion.models.ometa_classification import OMetaTagAndClassifica
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.source.database.database_service import DataModelLink
@@ -133,7 +133,7 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
     """
 
     topology = DbtServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     source_config: DbtPipeline
 
     @property
@@ -162,20 +162,22 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
     def get_dbt_files(self) -> Iterable[DbtFiles]:
         dbt_files = get_dbt_details(self.source_config.dbtConfigSource)
         for dbt_file in dbt_files:
-            self.context.dbt_file = dbt_file
+            self.context.get().dbt_file = dbt_file
             yield dbt_file
 
     def get_dbt_objects(self) -> Iterable[DbtObjects]:
         self.remove_manifest_non_required_keys(
-            manifest_dict=self.context.dbt_file.dbt_manifest
+            manifest_dict=self.context.get().dbt_file.dbt_manifest
         )
         dbt_objects = DbtObjects(
-            dbt_catalog=parse_catalog(self.context.dbt_file.dbt_catalog)
-            if self.context.dbt_file.dbt_catalog
+            dbt_catalog=parse_catalog(self.context.get().dbt_file.dbt_catalog)
+            if self.context.get().dbt_file.dbt_catalog
             else None,
-            dbt_manifest=parse_manifest(self.context.dbt_file.dbt_manifest),
-            dbt_run_results=parse_run_results(self.context.dbt_file.dbt_run_results)
-            if self.context.dbt_file.dbt_run_results
+            dbt_manifest=parse_manifest(self.context.get().dbt_file.dbt_manifest),
+            dbt_run_results=parse_run_results(
+                self.context.get().dbt_file.dbt_run_results
+            )
+            if self.context.get().dbt_file.dbt_run_results
             else None,
         )
         yield dbt_objects
@@ -204,7 +206,7 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
         """
         Prepare the data models
         """
-        for data_model_link in self.context.data_model_links:
+        for data_model_link in self.context.get().data_model_links:
             yield data_model_link
 
     @abstractmethod
@@ -231,7 +233,7 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
         """
         Prepare the DBT tests
         """
-        for _, dbt_test in self.context.dbt_tests.items():
+        for _, dbt_test in self.context.get().dbt_tests.items():
             yield dbt_test
 
     @abstractmethod