openmetadata-ingestion 1.6.3.2.dev1__py3-none-any.whl → 1.6.4.0.dev0__py3-none-any.whl

metadata/ingestion/source/dashboard/powerbi/client.py

@@ -47,7 +47,10 @@ from metadata.utils.logger import utils_logger
 
 logger = utils_logger()
 
-
+GETGROUPS_DEFAULT_PARAMS = {"$top": "1", "$skip": "0"}
+API_RESPONSE_MESSAGE_KEY = "message"
+AUTH_TOKEN_MAX_RETRIES = 5
+AUTH_TOKEN_RETRY_WAIT = 120
 # Similar inner methods with mode client. That's fine.
 # pylint: disable=duplicate-code
 class PowerBiApiClient:
@@ -59,6 +62,9 @@ class PowerBiApiClient:
 
     def __init__(self, config: PowerBIConnection):
         self.config = config
+        self.pagination_entity_per_page = min(
+            100, self.config.pagination_entity_per_page
+        )
         self.msal_client = msal.ConfidentialClientApplication(
             client_id=self.config.clientId,
             client_credential=self.config.clientSecret.get_secret_value(),
@@ -82,42 +88,84 @@ class PowerBiApiClient:
         """
         logger.info("Generating PowerBi access token")
 
-        response_data = self.msal_client.acquire_token_silent(
-            scopes=self.config.scope, account=None
-        )
-
+        response_data = self.get_auth_token_from_cache()
         if not response_data:
             logger.info("Token does not exist in the cache. Getting a new token.")
-            response_data = self.msal_client.acquire_token_for_client(
-                scopes=self.config.scope
-            )
+            response_data = self.generate_new_auth_token()
+        response_data = response_data or {}
         auth_response = PowerBiToken(**response_data)
         if not auth_response.access_token:
             raise InvalidSourceException(
-                "Failed to generate the PowerBi access token. Please check provided config"
+                f"Failed to generate the PowerBi access token. Please check provided config {response_data}"
             )
 
         logger.info("PowerBi Access Token generated successfully")
         return auth_response.access_token, auth_response.expires_in
 
+    def generate_new_auth_token(self) -> Optional[dict]:
+        """generate new auth token"""
+        retry = AUTH_TOKEN_MAX_RETRIES
+        while retry:
+            try:
+                response_data = self.msal_client.acquire_token_for_client(
+                    scopes=self.config.scope
+                )
+                return response_data
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error generating new auth token: {exc}")
+                # wait for time and retry
+                retry -= 1
+                if retry:
+                    logger.warning(
+                        f"Error generating new token: {exc}, "
+                        f"sleep {AUTH_TOKEN_RETRY_WAIT} seconds retrying {retry} more times.."
+                    )
+                    sleep(AUTH_TOKEN_RETRY_WAIT)
+                else:
+                    logger.warning(
+                        "Could not generate new token after maximum retries, "
+                        "Please check provided configs"
+                    )
+        return None
+
+    def get_auth_token_from_cache(self) -> Optional[dict]:
+        """fetch auth token from cache"""
+        retry = AUTH_TOKEN_MAX_RETRIES
+        while retry:
+            try:
+                response_data = self.msal_client.acquire_token_silent(
+                    scopes=self.config.scope, account=None
+                )
+                return response_data
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error getting token from cache: {exc}")
+                retry -= 1
+                if retry:
+                    logger.warning(
+                        f"Error getting token from cache: {exc}, "
+                        f"sleep {AUTH_TOKEN_RETRY_WAIT} seconds retrying {retry} more times.."
+                    )
+                    sleep(AUTH_TOKEN_RETRY_WAIT)
+                else:
+                    logger.warning(
+                        "Could not get token from cache after maximum retries, "
+                        "Please check provided configs"
+                    )
+        return None
+
     def fetch_dashboards(self) -> Optional[List[PowerBIDashboard]]:
         """Get dashboards method
         Returns:
             List[PowerBIDashboard]
         """
-        try:
-            if self.config.useAdminApis:
-                response_data = self.client.get("/myorg/admin/dashboards")
-                response = DashboardsResponse(**response_data)
-                return response.value
-            group = self.fetch_all_workspaces()[0]
-            return self.fetch_all_org_dashboards(group_id=group.id)
-
-        except Exception as exc:  # pylint: disable=broad-except
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Error fetching dashboards: {exc}")
-
-        return None
+        if self.config.useAdminApis:
+            response_data = self.client.get("/myorg/admin/dashboards")
+            response = DashboardsResponse(**response_data)
+            return response.value
+        group = self.fetch_all_workspaces()[0]
+        return self.fetch_all_org_dashboards(group_id=group.id)
 
     def fetch_all_org_dashboards(
         self, group_id: str
@@ -205,6 +253,7 @@ class PowerBiApiClient:
 
         return None
 
+    # pylint: disable=too-many-branches,too-many-statements
     def fetch_all_workspaces(self) -> Optional[List[Group]]:
         """Method to fetch all powerbi workspace details
         Returns:
@@ -213,28 +262,94 @@ class PowerBiApiClient:
         try:
             admin = "admin/" if self.config.useAdminApis else ""
             api_url = f"/myorg/{admin}groups"
-            entities_per_page = self.config.pagination_entity_per_page
-            params_data = {"$top": "1"}
-            response_data = self.client.get(api_url, data=params_data)
-            response = GroupsResponse(**response_data)
-            count = response.odata_count
+            entities_per_page = self.pagination_entity_per_page
+            failed_indexes = []
+            params_data = GETGROUPS_DEFAULT_PARAMS
+            response = self.client.get(api_url, data=params_data)
+            if (
+                not response
+                or API_RESPONSE_MESSAGE_KEY in response
+                or len(response) != len(GroupsResponse.__annotations__)
+            ):
+                logger.warning("Error fetching workspaces between results: (0, 1)")
+                if response and response.get(API_RESPONSE_MESSAGE_KEY):
+                    logger.warning(
+                        "Error message from API response: "
+                        f"{str(response.get(API_RESPONSE_MESSAGE_KEY))}"
+                    )
+                failed_indexes.append(params_data)
+                count = 0
+            else:
+                try:
+                    response = GroupsResponse(**response)
+                    count = response.odata_count
+                except Exception as exc:
+                    logger.warning(f"Error processing GetGroups response: {exc}")
+                    count = 0
             indexes = math.ceil(count / entities_per_page)
-
             workspaces = []
             for index in range(indexes):
                 params_data = {
                     "$top": str(entities_per_page),
                     "$skip": str(index * entities_per_page),
                 }
-                response_data = self.client.get(api_url, data=params_data)
-                if not response_data:
-                    logger.error(
-                        "Error fetching workspaces between results: "
-                        f"{str(index * entities_per_page)} - {str(entities_per_page)}"
+                response = self.client.get(api_url, data=params_data)
+                if (
+                    not response
+                    or API_RESPONSE_MESSAGE_KEY in response
+                    or len(response) != len(GroupsResponse.__annotations__)
+                ):
+                    index_range = (
+                        int(params_data.get("$skip")),
+                        int(params_data.get("$skip")) + int(params_data.get("$top")),
+                    )
+                    logger.warning(
+                        f"Error fetching workspaces between results: {str(index_range)}"
                     )
+                    if response and response.get(API_RESPONSE_MESSAGE_KEY):
+                        logger.warning(
+                            "Error message from API response: "
+                            f"{str(response.get(API_RESPONSE_MESSAGE_KEY))}"
+                        )
+                    failed_indexes.append(params_data)
                     continue
-                response = GroupsResponse(**response_data)
-                workspaces.extend(response.value)
+                try:
+                    response = GroupsResponse(**response)
+                    workspaces.extend(response.value)
+                except Exception as exc:
+                    logger.warning(f"Error processing GetGroups response: {exc}")
+
+            if failed_indexes:
+                logger.info(
+                    "Retrying one more time on failed indexes to get workspaces"
+                )
+                for params_data in failed_indexes:
+                    response = self.client.get(api_url, data=params_data)
+                    if (
+                        not response
+                        or API_RESPONSE_MESSAGE_KEY in response
+                        or len(response) != len(GroupsResponse.__annotations__)
+                    ):
+                        index_range = (
+                            int(params_data.get("$skip")),
+                            int(params_data.get("$skip"))
+                            + int(params_data.get("$top")),
+                        )
+                        logger.warning(
+                            f"Workspaces between results {str(index_range)} "
+                            "could not be fetched on multiple attempts"
+                        )
+                        if response and response.get(API_RESPONSE_MESSAGE_KEY):
+                            logger.warning(
+                                "Error message from API response: "
+                                f"{str(response.get(API_RESPONSE_MESSAGE_KEY))}"
+                            )
+                        continue
+                    try:
+                        response = GroupsResponse(**response)
+                        workspaces.extend(response.value)
+                    except Exception as exc:
+                        logger.warning(f"Error processing GetGroups response: {exc}")
             return workspaces
         except Exception as exc:  # pylint: disable=broad-except
             logger.debug(traceback.format_exc())

metadata/ingestion/source/database/bigquery/metadata.py

@@ -34,9 +34,11 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.data.table import (
+    ConstraintType,
     PartitionColumnDetails,
     PartitionIntervalTypes,
     Table,
+    TableConstraint,
     TablePartition,
     TableType,
 )
@@ -96,6 +98,7 @@ from metadata.ingestion.source.database.life_cycle_query_mixin import (
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.utils import fqn
 from metadata.utils.credentials import GOOGLE_CREDENTIALS
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.filters import filter_by_database, filter_by_schema
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
@@ -657,6 +660,42 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
             )
         return None
 
+    @calculate_execution_time()
+    def update_table_constraints(
+        self,
+        table_name,
+        schema_name,
+        db_name,
+        table_constraints,
+        foreign_columns,
+        columns,
+    ) -> List[TableConstraint]:
+        """
+        From topology.
+        process the table constraints of all tables
+        """
+        table_constraints = super().update_table_constraints(
+            table_name,
+            schema_name,
+            db_name,
+            table_constraints,
+            foreign_columns,
+            columns,
+        )
+        try:
+            table = self.client.get_table(fqn._build(db_name, schema_name, table_name))
+            if hasattr(table, "clustering_fields") and table.clustering_fields:
+                table_constraints.append(
+                    TableConstraint(
+                        constraintType=ConstraintType.CLUSTER_KEY,
+                        columns=table.clustering_fields,
+                    )
+                )
+        except Exception as exc:
+            logger.warning(f"Error getting clustering fields for {table_name}: {exc}")
+            logger.debug(traceback.format_exc())
+        return table_constraints
+
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
     ) -> Tuple[bool, Optional[TablePartition]]:

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

@@ -38,9 +38,9 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.databricks.client import DatabricksClient
 from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
+    DATABRICKS_SQL_STATEMENT_TEST,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
@@ -81,7 +81,6 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    client = DatabricksClient(service_connection)
 
     def test_database_query(engine: Engine, statement: str):
         """
@@ -106,7 +105,13 @@ def test_connection(
             engine=connection,
             statement=DATABRICKS_GET_CATALOGS,
         ),
-        "GetQueries": client.test_query_api_access,
+        "GetQueries": partial(
+            test_database_query,
+            engine=connection,
+            statement=DATABRICKS_SQL_STATEMENT_TEST.format(
+                query_history=service_connection.queryHistoryTable
+            ),
+        ),
     }
 
     return test_connection_steps(

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

@@ -11,12 +11,10 @@
 """
 Databricks lineage module
 """
-import traceback
-from datetime import datetime
-from typing import Iterator
 
-from metadata.generated.schema.type.basic import DateTime
-from metadata.generated.schema.type.tableQuery import TableQuery
+from metadata.ingestion.source.database.databricks.queries import (
+    DATABRICKS_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.databricks.query_parser import (
     DatabricksQueryParserSource,
 )
@@ -31,23 +29,13 @@ class DatabricksLineageSource(DatabricksQueryParserSource, LineageSource):
     Databricks Lineage Legacy Source
     """
 
-    def yield_table_query(self) -> Iterator[TableQuery]:
-        data = self.client.list_query_history(
-            start_date=self.start,
-            end_date=self.end,
+    sql_stmt = DATABRICKS_SQL_STATEMENT
+
+    filters = """
+        AND (
+            lower(statement_text) LIKE '%%create%%select%%'
+            OR lower(statement_text) LIKE '%%insert%%into%%select%%'
+            OR lower(statement_text) LIKE '%%update%%'
+            OR lower(statement_text) LIKE '%%merge%%'
         )
-        for row in data or []:
-            try:
-                if self.client.is_query_valid(row):
-                    yield TableQuery(
-                        dialect=self.dialect.value,
-                        query=row.get("query_text"),
-                        userName=row.get("user_name"),
-                        startTime=str(row.get("query_start_time_ms")),
-                        endTime=str(row.get("execution_end_time_ms")),
-                        analysisDate=DateTime(datetime.now()),
-                        serviceName=self.config.serviceName,
-                    )
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.warning(f"Error processing query_dict {row}: {exc}")
+    """

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

@@ -765,6 +765,8 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                 table_name=table_name,
                 schema=schema_name,
             )
+            found_location = False
+            log_location = None
             for result in list(cursor):
                 data = result.values()
                 if data[0] and data[0].strip() == "Comment":
@@ -777,6 +779,12 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
                         if data and data[1] and not data[1].startswith("dbfs")
                         else None
                     )
+                    found_location = True
+                    log_location = data[1]
+            if found_location and log_location:
+                logger.debug(f"Table {schema_name}.{table_name} has location {log_location}")
+            else:
+                logger.debug(f"Table {schema_name}.{table_name} does not have location")
 
         # Catch any exception without breaking the ingestion
         except Exception as exc:  # pylint: disable=broad-except

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

@@ -14,6 +14,30 @@ SQL Queries used during ingestion
 
 import textwrap
 
+DATABRICKS_SQL_STATEMENT = textwrap.dedent(
+    """
+    SELECT
+      statement_type AS query_type,
+      statement_text AS query_text,
+      executed_by AS user_name,
+      start_time AS start_time,
+      null AS database_name,
+      null AS schema_name,
+      end_time AS end_time,
+      total_duration_ms/1000 AS duration
+    from {query_history}
+    WHERE statement_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+    AND statement_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+    AND start_time between to_timestamp('{start_time}') and to_timestamp('{end_time}')
+    {filters}
+    LIMIT {result_limit}
+    """
+)
+
+DATABRICKS_SQL_STATEMENT_TEST = """
+ SELECT statement_text from  {query_history} LIMIT 1
+"""
+
 DATABRICKS_VIEW_DEFINITIONS = textwrap.dedent(
     """
     select

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

@@ -22,7 +22,6 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.databricks.client import DatabricksClient
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
 from metadata.utils.logger import ingestion_logger
 
@@ -36,18 +35,6 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
 
     filters: str
 
-    def _init_super(
-        self,
-        config: WorkflowSource,
-        metadata: OpenMetadata,
-    ):
-        super().__init__(config, metadata, False)
-
-    # pylint: disable=super-init-not-called
-    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
-        self._init_super(config=config, metadata=metadata)
-        self.client = DatabricksClient(self.service_connection)
-
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
@@ -61,7 +48,16 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
             )
         return cls(config, metadata)
 
-    def prepare(self):
+    def get_sql_statement(self, start_time, end_time):
         """
-        By default, there's nothing to prepare
+        returns sql statement to fetch query logs.
+
+        Override if we have specific parameters
         """
+        return self.sql_stmt.format(
+            start_time=start_time,
+            end_time=end_time,
+            filters=self.get_filters(),
+            result_limit=self.source_config.resultLimit,
+            query_history=self.service_connection.queryHistoryTable,
+        )

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

@@ -11,12 +11,10 @@
 """
 Databricks usage module
 """
-import traceback
-from datetime import datetime
-from typing import Iterable
 
-from metadata.generated.schema.type.basic import DateTime
-from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.ingestion.source.database.databricks.queries import (
+    DATABRICKS_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.databricks.query_parser import (
     DatabricksQueryParserSource,
 )
@@ -31,36 +29,8 @@ class DatabricksUsageSource(DatabricksQueryParserSource, UsageSource):
     Databricks Usage Source
     """
 
-    def yield_table_queries(self) -> Iterable[TableQuery]:
-        """
-        Method to yield TableQueries
-        """
-        queries = []
-        data = self.client.list_query_history(
-            start_date=self.start,
-            end_date=self.end,
-        )
-        for row in data or []:
-            try:
-                if self.client.is_query_valid(row):
-                    queries.append(
-                        TableQuery(
-                            dialect=self.dialect.value,
-                            query=row.get("query_text"),
-                            userName=row.get("user_name"),
-                            startTime=str(row.get("query_start_time_ms")),
-                            endTime=str(row.get("execution_end_time_ms")),
-                            analysisDate=DateTime(datetime.now()),
-                            serviceName=self.config.serviceName,
-                            duration=row.get("duration")
-                            if row.get("duration")
-                            else None,
-                        )
-                    )
-            except Exception as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Failed to process query {row.get('query_text')} due to: {err}"
-                )
+    sql_stmt = DATABRICKS_SQL_STATEMENT
 
-        yield TableQueries(queries=queries)
+    filters = """
+        AND statement_type NOT IN ('SHOW', 'DESCRIBE', 'USE')
+    """

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

@@ -50,7 +50,7 @@ def get_connection(connection: Db2Connection) -> Engine:
             "w",
             encoding=UTF_8,
         ) as file:
-            file.write(connection.license)
+            file.write(connection.license.encode(UTF_8).decode("unicode-escape"))
 
     return create_generic_db_connection(
         connection=connection,

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

@@ -1,4 +1,3 @@
-from metadata.ingestion.source.database.db2.lineage import Db2LineageSource
 from metadata.ingestion.source.database.db2.metadata import Db2Source
 from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
     DB2ProfilerInterface,
@@ -6,7 +5,5 @@ from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=Db2Source,
-    profiler_class=DB2ProfilerInterface,
-    lineage_source_class=Db2LineageSource,
+    metadata_source_class=Db2Source, profiler_class=DB2ProfilerInterface
 )

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

@@ -1093,7 +1093,9 @@ class DbtSource(DbtServiceSource):
 
                 # Create the test case result object
                 test_case_result = TestCaseResult(
-                    timestamp=Timestamp(datetime_to_timestamp(dbt_timestamp)),
+                    timestamp=Timestamp(
+                        datetime_to_timestamp(dbt_timestamp, milliseconds=True)
+                    ),
                     testCaseStatus=test_case_status,
                     testResultValue=[
                         TestResultValue(

metadata/ingestion/source/database/druid/service_spec.py

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.druid.lineage import DruidLineageSource
 from metadata.ingestion.source.database.druid.metadata import DruidSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=DruidSource,
-    lineage_source_class=DruidLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=DruidSource)

metadata/ingestion/source/database/greenplum/service_spec.py

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.greenplum.lineage import GreenplumLineageSource
 from metadata.ingestion.source.database.greenplum.metadata import GreenplumSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=GreenplumSource,
-    lineage_source_class=GreenplumLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=GreenplumSource)

metadata/ingestion/source/database/hive/service_spec.py

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.hive.lineage import HiveLineageSource
 from metadata.ingestion.source.database.hive.metadata import HiveSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=HiveSource,
-    lineage_source_class=HiveLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=HiveSource)

metadata/ingestion/source/database/impala/service_spec.py

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.impala.lineage import ImpalaLineageSource
 from metadata.ingestion.source.database.impala.metadata import ImpalaSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=ImpalaSource,
-    lineage_source_class=ImpalaLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=ImpalaSource)

metadata/ingestion/source/database/lineage_source.py

@@ -26,7 +26,10 @@ from metadata.generated.schema.type.basic import FullyQualifiedEntityName, SqlQu
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper, Dialect
-from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
+from metadata.ingestion.lineage.sql_lineage import (
+    get_lineage_by_graph,
+    get_lineage_by_query,
+)
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
@@ -195,6 +198,12 @@ class LineageSource(QueryParserSource, ABC):
     def query_lineage_generator(
         self, table_queries: List[TableQuery], queue: Queue
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+        if self.graph is None and self.source_config.enableTempTableLineage:
+            import networkx as nx
+
+            # Create a directed graph
+            self.graph = nx.DiGraph()
+
         for table_query in table_queries or []:
             if not self._query_already_processed(table_query):
                 lineages: Iterable[Either[AddLineageRequest]] = get_lineage_by_query(
@@ -205,6 +214,7 @@ class LineageSource(QueryParserSource, ABC):
                     schema_name=table_query.databaseSchema,
                     dialect=self.dialect,
                     timeout_seconds=self.source_config.parsingTimeoutLimit,
+                    graph=self.graph,
                 )
 
                 for lineage_request in lineages or []:
@@ -298,6 +308,7 @@ class LineageSource(QueryParserSource, ABC):
         if self.source_config.processQueryLineage:
             if hasattr(self.service_connection, "supportsLineageExtraction"):
                 yield from self.yield_query_lineage() or []
+                yield from get_lineage_by_graph(graph=self.graph)
             else:
                 logger.warning(
                     f"Lineage extraction is not supported for {str(self.service_connection.type.value)} connection"