openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0__py3-none-any.whl

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

@@ -45,41 +45,155 @@ SNOWFLAKE_FETCH_ALL_TAGS = textwrap.dedent(
 """
 )
 
-SNOWFLAKE_GET_TABLE_NAMES = """
-select TABLE_NAME from information_schema.tables 
-where TABLE_SCHEMA = '{}' and TABLE_TYPE = 'BASE TABLE'
-"""
-
 SNOWFLAKE_GET_EXTERNAL_TABLE_NAMES = """
-select TABLE_NAME from information_schema.tables 
-where TABLE_SCHEMA = '{}' AND TABLE_TYPE = 'EXTERNAL TABLE'
+select TABLE_NAME, NULL from information_schema.tables
+where TABLE_SCHEMA = '{schema}' AND TABLE_TYPE = 'EXTERNAL TABLE'
 """
 
+SNOWFLAKE_INCREMENTAL_GET_EXTERNAL_TABLE_NAMES = """
+select TABLE_NAME, DELETED
+from (
+    select
+        TABLE_NAME,
+        DELETED,
+        ROW_NUMBER() over (
+            partition by TABLE_NAME order by LAST_DDL desc
+        ) as ROW_NUMBER
+    from snowflake.account_usage.tables
+    where TABLE_CATALOG = '{database}'
+      and TABLE_SCHEMA = '{schema}'
+      and TABLE_TYPE = 'EXTERNAL TABLE'
+      and DATE_PART(epoch_millisecond, LAST_DDL) >= '{date}'
+)
+where ROW_NUMBER = 1
+"""
 
 SNOWFLAKE_GET_WITHOUT_TRANSIENT_TABLE_NAMES = """
-select TABLE_NAME from information_schema.tables 
-where TABLE_SCHEMA = '{}' 
-AND TABLE_TYPE = 'BASE TABLE' 
+select TABLE_NAME, NULL from information_schema.tables
+where TABLE_SCHEMA = '{schema}'
+AND TABLE_TYPE = 'BASE TABLE'
 AND IS_TRANSIENT != 'YES'
+AND IS_DYNAMIC != 'YES'
+"""
+
+SNOWFLAKE_INCREMENTAL_GET_WITHOUT_TRANSIENT_TABLE_NAMES = """
+select TABLE_NAME, DELETED
+from (
+    select
+        TABLE_NAME,
+        DELETED,
+        ROW_NUMBER() over (
+            partition by TABLE_NAME order by LAST_DDL desc
+        ) as ROW_NUMBER
+    from snowflake.account_usage.tables
+    where TABLE_CATALOG = '{database}'
+    and TABLE_SCHEMA = '{schema}'
+    and TABLE_TYPE = 'BASE TABLE'
+    and IS_TRANSIENT != 'YES'
+    AND IS_DYNAMIC != 'YES'
+    and DATE_PART(epoch_millisecond, LAST_DDL) >= '{date}'
+)
+where ROW_NUMBER = 1
 """
 
 SNOWFLAKE_GET_VIEW_NAMES = """
-select TABLE_NAME from information_schema.tables 
-where TABLE_SCHEMA = '{}' and TABLE_TYPE = 'VIEW'
+select TABLE_NAME, NULL from information_schema.tables
+where TABLE_SCHEMA = '{schema}' and TABLE_TYPE = 'VIEW'
+"""
+
+SNOWFLAKE_INCREMENTAL_GET_VIEW_NAMES = """
+select TABLE_NAME, DELETED
+from (
+    select
+        TABLE_NAME,
+        DELETED,
+        ROW_NUMBER() over (
+            partition by TABLE_NAME order by LAST_DDL desc
+        ) as ROW_NUMBER
+    from snowflake.account_usage.tables
+    where  TABLE_CATALOG = '{database}'
+    and TABLE_SCHEMA = '{schema}'
+    and TABLE_TYPE = 'VIEW'
+    and DATE_PART(epoch_millisecond, LAST_DDL) >= '{date}'
+)
+where ROW_NUMBER = 1
 """
 
 SNOWFLAKE_GET_MVIEW_NAMES = """
-select TABLE_NAME from information_schema.tables 
-where TABLE_SCHEMA = '{}' and TABLE_TYPE = 'MATERIALIZED VIEW'
+select TABLE_NAME, NULL from information_schema.tables
+where TABLE_SCHEMA = '{schema}' and TABLE_TYPE = 'MATERIALIZED VIEW'
+"""
+
+SNOWFLAKE_INCREMENTAL_GET_MVIEW_NAMES = """
+select TABLE_NAME, DELETED
+from (
+    select
+        TABLE_NAME,
+        DELETED,
+        ROW_NUMBER() over (
+            partition by TABLE_NAME order by LAST_DDL desc
+        ) as ROW_NUMBER
+    from snowflake.account_usage.tables
+    where  TABLE_CATALOG = '{database}'
+    and TABLE_SCHEMA = '{schema}'
+    and TABLE_TYPE = 'MATERIALIZED VIEW'
+    and DATE_PART(epoch_millisecond, LAST_DDL) >= '{date}'
+)
+where ROW_NUMBER = 1
 """
 
 SNOWFLAKE_GET_TRANSIENT_NAMES = """
-select TABLE_NAME from information_schema.tables 
-where TABLE_SCHEMA = '{}' 
-AND TABLE_TYPE = 'BASE TABLE' 
+select TABLE_NAME, NULL from information_schema.tables
+where TABLE_SCHEMA = '{schema}'
+AND TABLE_TYPE = 'BASE TABLE'
 AND IS_TRANSIENT = 'YES'
 """
 
+SNOWFLAKE_INCREMENTAL_GET_TRANSIENT_NAMES = """
+select TABLE_NAME, DELETED
+from (
+    select
+        TABLE_NAME,
+        DELETED,
+        ROW_NUMBER() over (
+            partition by TABLE_NAME order by LAST_DDL desc
+        ) as ROW_NUMBER
+    from snowflake.account_usage.tables
+    where TABLE_CATALOG = '{database}'
+    and TABLE_SCHEMA = '{schema}'
+    and TABLE_TYPE = 'BASE TABLE'
+    and IS_TRANSIENT = 'YES'
+    and DATE_PART(epoch_millisecond, LAST_DDL) >= '{date}'
+)
+where ROW_NUMBER = 1
+"""
+
+SNOWFLAKE_GET_DYNAMIC_TABLE_NAMES = """
+select TABLE_NAME, NULL from information_schema.tables
+where TABLE_SCHEMA = '{schema}'
+AND TABLE_TYPE = 'BASE TABLE'
+AND IS_DYNAMIC = 'YES'
+"""
+
+SNOWFLAKE_INCREMENTAL_GET_DYNAMIC_TABLE_NAMES = """
+select TABLE_NAME, DELETED
+from (
+    select
+        TABLE_NAME,
+        DELETED,
+        ROW_NUMBER() over (
+            partition by TABLE_NAME order by LAST_DDL desc
+        ) as ROW_NUMBER
+    from snowflake.account_usage.tables
+    where TABLE_CATALOG = '{database}'
+    and TABLE_SCHEMA = '{schema}'
+    and TABLE_TYPE = 'BASE TABLE'
+    and IS_DYNAMIC = 'YES'
+    and DATE_PART(epoch_millisecond, LAST_DDL) >= '{date}'
+)
+where ROW_NUMBER = 1
+"""
+
 SNOWFLAKE_GET_COMMENTS = textwrap.dedent(
     """
   select
@@ -96,17 +210,17 @@ SNOWFLAKE_GET_CLUSTER_KEY = """
   select CLUSTERING_KEY,
           TABLE_SCHEMA,
           TABLE_NAME
-  from   information_schema.tables 
+  from   information_schema.tables
   where  TABLE_TYPE = 'BASE TABLE'
   and CLUSTERING_KEY is not null
 """
 
 
 SNOWFLAKE_GET_SCHEMA_COMMENTS = """
-SELECT 
+SELECT
       catalog_name DATABASE_NAME,
       SCHEMA_NAME,
-      COMMENT 
+      COMMENT
 FROM information_schema.schemata
 """
 
@@ -115,6 +229,10 @@ SNOWFLAKE_GET_DATABASE_COMMENTS = """
 select DATABASE_NAME,COMMENT from information_schema.databases
 """
 
+SNOWFLAKE_GET_EXTERNAL_LOCATIONS = """
+SHOW EXTERNAL TABLES IN DATABASE "{database_name}"
+"""
+
 SNOWFLAKE_TEST_FETCH_TAG = """
 select TAG_NAME from snowflake.account_usage.tag_references limit 1
 """
@@ -124,7 +242,7 @@ SELECT query_text from snowflake.account_usage.query_history limit 1
 """
 
 SNOWFLAKE_TEST_GET_TABLES = """
-SELECT TABLE_NAME FROM "{database_name}".information_schema.tables LIMIT 1 
+SELECT TABLE_NAME FROM "{database_name}".information_schema.tables LIMIT 1
 """
 
 SNOWFLAKE_GET_DATABASES = "SHOW DATABASES"
@@ -155,7 +273,7 @@ SNOWFLAKE_GET_CURRENT_ACCOUNT = "SELECT CURRENT_ACCOUNT_NAME() AS ACCOUNT"
 
 SNOWFLAKE_LIFE_CYCLE_QUERY = textwrap.dedent(
     """
-select 
+select
 table_name as table_name,
 created as created_at
 from snowflake.account_usage.tables
@@ -166,7 +284,7 @@ and table_catalog = '{database_name}'
 
 SNOWFLAKE_GET_STORED_PROCEDURES = textwrap.dedent(
     """
-SELECT 
+SELECT
   PROCEDURE_NAME AS name,
   PROCEDURE_OWNER AS owner,
   PROCEDURE_LANGUAGE AS language,
@@ -193,7 +311,7 @@ WITH SP_HISTORY AS (
       END_TIME
     FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY SP
     WHERE QUERY_TYPE = 'CALL'
-      AND START_TIME >= '{start_date}' 
+      AND START_TIME >= '{start_date}'
 ),
 Q_HISTORY AS (
     SELECT
@@ -210,7 +328,7 @@ Q_HISTORY AS (
     WHERE QUERY_TYPE <> 'CALL'
       AND QUERY_TEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
       AND QUERY_TEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'
-      AND START_TIME >= '{start_date}' 
+      AND START_TIME >= '{start_date}'
 )
 SELECT
   Q.QUERY_TYPE AS QUERY_TYPE,
@@ -233,3 +351,7 @@ JOIN Q_HISTORY Q
 ORDER BY PROCEDURE_START_TIME DESC
     """
 )
+
+SNOWFLAKE_GET_TABLE_DDL = """
+SELECT GET_DDL('TABLE','{table_name}') AS \"text\"
+"""

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

@@ -40,7 +40,9 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
     """
 
     @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: SnowflakeConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SnowflakeConnection):

metadata/ingestion/source/database/snowflake/utils.py

@@ -13,7 +13,7 @@
 Module to define overriden dialect methods
 """
 
-from typing import Optional
+from typing import Dict, Optional
 
 import sqlalchemy.types as sqltypes
 from sqlalchemy import exc as sa_exc
@@ -22,14 +22,29 @@ from sqlalchemy.engine import reflection
 from sqlalchemy.sql import text
 from sqlalchemy.types import FLOAT
 
+from metadata.ingestion.source.database.incremental_metadata_extraction import (
+    IncrementalConfig,
+)
+from metadata.ingestion.source.database.snowflake.models import (
+    SnowflakeTable,
+    SnowflakeTableList,
+)
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_GET_COMMENTS,
+    SNOWFLAKE_GET_DYNAMIC_TABLE_NAMES,
     SNOWFLAKE_GET_EXTERNAL_TABLE_NAMES,
     SNOWFLAKE_GET_MVIEW_NAMES,
     SNOWFLAKE_GET_SCHEMA_COLUMNS,
+    SNOWFLAKE_GET_TABLE_DDL,
     SNOWFLAKE_GET_TRANSIENT_NAMES,
     SNOWFLAKE_GET_VIEW_NAMES,
     SNOWFLAKE_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_DYNAMIC_TABLE_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_EXTERNAL_TABLE_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_MVIEW_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_TRANSIENT_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_VIEW_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
 )
 from metadata.utils import fqn
 from metadata.utils.sqlalchemy_utils import (
@@ -37,6 +52,36 @@ from metadata.utils.sqlalchemy_utils import (
     get_table_comment_wrapper,
 )
 
+Query = str
+QueryMap = Dict[str, Query]
+
+
+TABLE_QUERY_MAPS = {
+    "full": {
+        "default": SNOWFLAKE_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
+        "transient_tables": SNOWFLAKE_GET_TRANSIENT_NAMES,
+        "external_tables": SNOWFLAKE_GET_EXTERNAL_TABLE_NAMES,
+        "dynamic_tables": SNOWFLAKE_GET_DYNAMIC_TABLE_NAMES,
+    },
+    "incremental": {
+        "default": SNOWFLAKE_INCREMENTAL_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
+        "transient_tables": SNOWFLAKE_INCREMENTAL_GET_TRANSIENT_NAMES,
+        "external_tables": SNOWFLAKE_INCREMENTAL_GET_EXTERNAL_TABLE_NAMES,
+        "dynamic_tables": SNOWFLAKE_INCREMENTAL_GET_DYNAMIC_TABLE_NAMES,
+    },
+}
+
+VIEW_QUERY_MAPS = {
+    "full": {
+        "views": SNOWFLAKE_GET_VIEW_NAMES,
+        "materialized_views": SNOWFLAKE_GET_MVIEW_NAMES,
+    },
+    "incremental": {
+        "views": SNOWFLAKE_INCREMENTAL_GET_VIEW_NAMES,
+        "materialized_views": SNOWFLAKE_INCREMENTAL_GET_MVIEW_NAMES,
+    },
+}
+
 
 def _quoted_name(entity_name: Optional[str]) -> Optional[str]:
     if entity_name:
@@ -88,25 +133,78 @@ def get_view_names_reflection(self, schema=None, **kw):
         )
 
 
-def get_table_names(self, connection, schema, **kw):
-    query = SNOWFLAKE_GET_WITHOUT_TRANSIENT_TABLE_NAMES
+def _get_query_map(
+    incremental: Optional[IncrementalConfig], query_maps: Dict[str, QueryMap]
+):
+    """Returns the proper queries depending if the extraction is Incremental or Full."""
+    if incremental and incremental.enabled:
+        return query_maps["incremental"]
+    return query_maps["full"]
+
+
+def _get_query_parameters(
+    self, connection, schema: str, incremental: Optional[IncrementalConfig]
+):
+    """Returns the proper query parameters depending if the extraciton is Incremental or Full"""
+    parameters = {"schema": fqn.unquote_name(schema)}
+
+    if incremental and incremental.enabled:
+        database, _ = self._current_database_schema(connection)  # pylint: disable=W0212
+        parameters = {
+            **parameters,
+            "date": incremental.start_timestamp,
+            "database": database,
+        }
+
+    return parameters
+
+
+def get_table_names(self, connection, schema: str, **kw):
+    """Return the Table names to process based on the incremental setup."""
+    incremental = kw.get("incremental")
+
+    queries = _get_query_map(incremental, TABLE_QUERY_MAPS)
+    parameters = _get_query_parameters(self, connection, schema, incremental)
+
+    query = queries["default"]
+
     if kw.get("include_transient_tables"):
-        query = SNOWFLAKE_GET_TRANSIENT_NAMES
+        query = queries["transient_tables"]
 
     if kw.get("external_tables"):
-        query = SNOWFLAKE_GET_EXTERNAL_TABLE_NAMES
-    cursor = connection.execute(query.format(fqn.unquote_name(schema)))
-    result = [self.normalize_name(row[0]) for row in cursor]
+        query = queries["external_tables"]
+
+    if kw.get("dynamic_tables"):
+        query = queries["dynamic_tables"]
+
+    cursor = connection.execute(query.format(**parameters))
+    result = SnowflakeTableList(
+        tables=[
+            SnowflakeTable(name=self.normalize_name(row[0]), deleted=row[1])
+            for row in cursor
+        ]
+    )
     return result
 
 
 def get_view_names(self, connection, schema, **kw):
+    incremental = kw.get("incremental")
+
+    queries = _get_query_map(incremental, VIEW_QUERY_MAPS)
+    parameters = _get_query_parameters(self, connection, schema, incremental)
+
     if kw.get("materialized_views"):
-        query = SNOWFLAKE_GET_MVIEW_NAMES
+        query = queries["materialized_views"]
     else:
-        query = SNOWFLAKE_GET_VIEW_NAMES
-    cursor = connection.execute(query.format(schema))
-    result = [self.normalize_name(row[0]) for row in cursor]
+        query = queries["views"]
+
+    cursor = connection.execute(query.format(**parameters))
+    result = SnowflakeTableList(
+        tables=[
+            SnowflakeTable(name=self.normalize_name(row[0]), deleted=row[1])
+            for row in cursor
+        ]
+    )
     return result
 
 
@@ -290,6 +388,65 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
     return foreign_key_map.get(table_name, [])
 
 
+@reflection.cache
+def get_schema_foreign_keys(self, connection, schema, **kw):
+    current_database, current_schema = self._current_database_schema(connection, **kw)
+    result = connection.execute(
+        text(
+            f"SHOW /* sqlalchemy:_get_schema_foreign_keys */ IMPORTED KEYS IN SCHEMA {schema}"
+        )
+    )
+    foreign_key_map = {}
+    for row in result:
+        name = self.normalize_name(row._mapping["fk_name"])
+        if name not in foreign_key_map:
+            referred_schema = self.normalize_name(row._mapping["pk_schema_name"])
+            foreign_key_map[name] = {
+                "constrained_columns": [
+                    self.normalize_name(row._mapping["fk_column_name"])
+                ],
+                # referred schema should be None in context where it doesn't need to be specified
+                # https://docs.sqlalchemy.org/en/14/core/reflection.html#reflection-schema-qualified-interaction
+                "referred_schema": (
+                    referred_schema
+                    if referred_schema not in (self.default_schema_name, current_schema)
+                    else None
+                ),
+                "referred_table": self.normalize_name(row._mapping["pk_table_name"]),
+                "referred_columns": [
+                    self.normalize_name(row._mapping["pk_column_name"])
+                ],
+                "referred_database": self.normalize_name(
+                    row._mapping["pk_database_name"]
+                ),
+                "name": name,
+                "table_name": self.normalize_name(row._mapping["fk_table_name"]),
+            }
+            options = {}
+            if self.normalize_name(row._mapping["delete_rule"]) != "NO ACTION":
+                options["ondelete"] = self.normalize_name(row._mapping["delete_rule"])
+            if self.normalize_name(row._mapping["update_rule"]) != "NO ACTION":
+                options["onupdate"] = self.normalize_name(row._mapping["update_rule"])
+            foreign_key_map[name]["options"] = options
+        else:
+            foreign_key_map[name]["constrained_columns"].append(
+                self.normalize_name(row._mapping["fk_column_name"])
+            )
+            foreign_key_map[name]["referred_columns"].append(
+                self.normalize_name(row._mapping["pk_column_name"])
+            )
+
+    ans = {}
+
+    for _, v in foreign_key_map.items():
+        if v["table_name"] not in ans:
+            ans[v["table_name"]] = []
+        ans[v["table_name"]].append(
+            {k2: v2 for k2, v2 in v.items() if k2 != "table_name"}
+        )
+    return ans
+
+
 @reflection.cache
 def get_unique_constraints(self, connection, table_name, schema, **kw):
     schema = schema or self.default_schema_name
@@ -320,3 +477,22 @@ def get_columns(self, connection, table_name, schema=None, **kw):
     if normalized_table_name not in schema_columns:
         raise sa_exc.NoSuchTableError()
     return schema_columns[normalized_table_name]
+
+
+@reflection.cache
+def get_table_ddl(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    Gets the Table DDL
+    """
+    schema = schema or self.default_schema_name
+    table_name = f"{schema}.{table_name}" if schema else table_name
+    cursor = connection.execute(SNOWFLAKE_GET_TABLE_DDL.format(table_name=table_name))
+    try:
+        result = cursor.fetchone()
+        if result:
+            return result[0]
+    except Exception:
+        pass
+    return None

metadata/ingestion/source/database/sql_column_handler.py

@@ -26,6 +26,7 @@ from metadata.generated.schema.entity.data.table import (
     TableConstraint,
 )
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.helpers import clean_up_starting_ending_double_quotes_in_string
 from metadata.utils.logger import ingestion_logger
 
@@ -195,6 +196,7 @@ class SqlColumnHandlerMixin:
             ]
         return Column(**parsed_string)
 
+    @calculate_execution_time()
     def get_columns_and_constraints(  # pylint: disable=too-many-locals
         self, schema_name: str, table_name: str, db_name: str, inspector: Inspector
     ) -> Tuple[

metadata/ingestion/source/database/sqlalchemy_source.py

@@ -24,7 +24,7 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.models.topology import TopologyContext
+from metadata.ingestion.models.topology import TopologyContextManager
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.logger import ingestion_logger
 
@@ -38,7 +38,7 @@ class SqlAlchemySource(ABC):
 
     engine: Engine
     metadata: OpenMetadata
-    context: TopologyContext
+    context: TopologyContextManager
     database_source_state: Set
     source_config: DatabaseServiceMetadataPipeline
     config: WorkflowSource
@@ -74,11 +74,11 @@ class SqlAlchemySource(ABC):
         """
 
     @abstractmethod
-    def get_view_definition(
+    def get_schema_definition(
         self, table_type, table_name: str, schema_name: str, inspector: Inspector
     ) -> Optional[str]:
         """
-        Method to fetch view definition
+        Method to fetch schema definition
         """
 
     @abstractmethod

metadata/ingestion/source/database/sqlite/metadata.py

@@ -14,6 +14,8 @@ Sqlite source implementation.
 Useful for testing!
 """
 
+from typing import Optional
+
 from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
     SQLiteConnection,
 )
@@ -32,7 +34,9 @@ class SqliteSource(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 = config.serviceConnection.__root__.config
         if not isinstance(connection, SQLiteConnection):

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -37,7 +37,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
-from metadata.ingestion.models.topology import TopologyContext
+from metadata.ingestion.models.topology import TopologyContextManager
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.stored_procedures import get_procedure_name_from_call
@@ -81,7 +81,7 @@ class StoredProcedureMixin(ABC):
     It should be inherited in those Sources that implement Stored Procedure ingestion.
     """
 
-    context: TopologyContext
+    context: TopologyContextManager
     status: Status
     source_config: DatabaseServiceMetadataPipeline
     engine: Engine
@@ -148,16 +148,16 @@ class StoredProcedureMixin(ABC):
         self, query_by_procedure: QueryByProcedure, procedure: StoredProcedure
     ) -> Iterable[Either[AddLineageRequest]]:
         """Add procedure lineage from its query"""
-        self.context.stored_procedure_query_lineage = False
+        self.context.get().stored_procedure_query_lineage = False
         if self.is_lineage_query(
             query_type=query_by_procedure.query_type,
             query_text=query_by_procedure.query_text,
         ):
-            self.context.stored_procedure_query_lineage = True
+            self.context.get().stored_procedure_query_lineage = True
             for either_lineage in get_lineage_by_query(
                 self.metadata,
                 query=query_by_procedure.query_text,
-                service_name=self.context.database_service,
+                service_name=self.context.get().database_service,
                 database_name=query_by_procedure.query_database_name,
                 schema_name=query_by_procedure.query_schema_name,
                 dialect=ConnectionTypeDialectMapper.dialect_of(
@@ -193,8 +193,10 @@ class StoredProcedureMixin(ABC):
                     id=procedure.id,
                     type="storedProcedure",
                 ),
-                processedLineage=bool(self.context.stored_procedure_query_lineage),
-                service=self.context.database_service,
+                processedLineage=bool(
+                    self.context.get().stored_procedure_query_lineage
+                ),
+                service=self.context.get().database_service,
             )
         )
 
@@ -202,12 +204,12 @@ class StoredProcedureMixin(ABC):
         self,
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
         """Get all the queries and procedures list and yield them"""
-        if self.context.stored_procedures:
+        if self.context.get().stored_procedures:
             logger.info("Processing Lineage for Stored Procedures")
             # First, get all the query history
             queries_dict = self.get_stored_procedure_queries_dict()
             # Then for each procedure, iterate over all its queries
-            for procedure_fqn in self.context.stored_procedures:
+            for procedure_fqn in self.context.get().stored_procedures:
                 procedure = self.metadata.get_by_name(
                     entity=StoredProcedure, fqn=procedure_fqn
                 )