openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

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

@@ -11,11 +11,22 @@
 """
 Snowflake usage module
 """
+import traceback
+from datetime import timedelta
+from typing import Iterable
+
+from metadata.generated.schema.type.basic import DateTime
+from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.ingestion.lineage.masker import mask_query
 from metadata.ingestion.source.database.snowflake.queries import SNOWFLAKE_SQL_STATEMENT
 from metadata.ingestion.source.database.snowflake.query_parser import (
+    SNOWFLAKE_QUERY_BATCH_SIZE,
     SnowflakeQueryParserSource,
 )
 from metadata.ingestion.source.database.usage_source import UsageSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class SnowflakeUsageSource(SnowflakeQueryParserSource, UsageSource):
@@ -41,3 +52,89 @@ class SnowflakeUsageSource(SnowflakeQueryParserSource, UsageSource):
         "INSERT",
         "MERGE",
     ]
+
+    def yield_table_queries(self) -> Iterable[TableQuery]:
+        """
+        Given an Engine, iterate over the day range and
+        query the results with pagination
+        """
+        daydiff = self.end - self.start
+        for days in range(daydiff.days):
+            logger.info(
+                f"Scanning query logs for {(self.start + timedelta(days=days)).date()} - "
+                f"{(self.start + timedelta(days=days + 1)).date()}"
+            )
+            query = None
+            offset = 0
+            total_fetched = 0
+            max_results = self.source_config.resultLimit
+            try:
+                for engine in self.get_engine():
+                    while total_fetched < max_results:
+                        batch_size = min(
+                            SNOWFLAKE_QUERY_BATCH_SIZE, max_results - total_fetched
+                        )
+                        query = self.get_sql_statement(
+                            start_time=self.start + timedelta(days=days),
+                            end_time=self.start + timedelta(days=days + 1),
+                            offset=offset,
+                            limit=batch_size,
+                        )
+                        with engine.connect() as conn:
+                            rows = conn.execute(query)
+                            queries = []
+                            row_count = 0
+                            for row in rows:
+                                row = dict(row)
+                                row_count += 1
+                                try:
+                                    row.update({k.lower(): v for k, v in row.items()})
+                                    logger.debug(f"Processing row: {row}")
+                                    query_type = row.get("query_type")
+                                    query_text = self.format_query(row["query_text"])
+                                    queries.append(
+                                        TableQuery(
+                                            query=query_text,
+                                            query_type=query_type,
+                                            exclude_usage=self.check_life_cycle_query(
+                                                query_type=query_type,
+                                                query_text=query_text,
+                                            ),
+                                            dialect=self.dialect.value,
+                                            userName=row["user_name"],
+                                            startTime=str(row["start_time"]),
+                                            endTime=str(row["end_time"]),
+                                            analysisDate=DateTime(row["start_time"]),
+                                            aborted=self.get_aborted_status(row),
+                                            databaseName=self.get_database_name(row),
+                                            duration=row.get("duration"),
+                                            serviceName=self.config.serviceName,
+                                            databaseSchema=self.get_schema_name(row),
+                                            cost=row.get("cost"),
+                                        )
+                                    )
+                                except Exception as exc:
+                                    logger.debug(traceback.format_exc())
+                                    logger.warning(
+                                        f"Unexpected exception processing row [{row}]: {exc}"
+                                    )
+                        if queries:
+                            yield TableQueries(queries=queries)
+                        total_fetched += row_count
+                        if row_count < batch_size:
+                            break
+                        offset += batch_size
+                        logger.info(
+                            f"Fetching next page with offset {offset} (fetched {total_fetched}/{max_results}) "
+                            f"for {(self.start + timedelta(days=days)).date()}"
+                        )
+            except Exception as exc:
+                if query:
+                    logger.debug(
+                        (
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value) or query}"
+                            "\n##########################"
+                        )
+                    )
+                logger.debug(traceback.format_exc())
+                logger.error(f"Source usage processing error: {exc}")

metadata/ingestion/source/database/sql_column_handler.py

@@ -268,9 +268,17 @@ class SqlColumnHandlerMixin:
 
         table_columns = []
 
-        columns = self._get_columns_internal(
-            schema_name, table_name, db_name, inspector, table_type
-        )
+        try:
+            columns = self._get_columns_internal(
+                schema_name, table_name, db_name, inspector, table_type
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unexpected exception getting columns for table [{table_name}] "
+                f"(schema: '{schema_name}', db: '{db_name}'): {exc}"
+            )
+            columns = []
 
         def process_column(column: dict):
             (

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -12,15 +12,11 @@
 Mixin class with common Stored Procedures logic aimed at lineage.
 """
 import json
-import re
 import traceback
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from datetime import datetime
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Iterator, Union
 
-from networkx import DiGraph
-from pydantic import BaseModel, ConfigDict, Field
 from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
@@ -32,16 +28,15 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.databaseServiceQueryLineagePipeline import (
     DatabaseServiceQueryLineagePipeline,
 )
-from metadata.generated.schema.type.basic import SqlQuery, Timestamp
-from metadata.generated.schema.type.entityLineage import Source as LineageSource
-from metadata.generated.schema.type.entityReference import EntityReference
 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.ometa_lineage import OMetaLineageRequest
-from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.lineage_processors import (
+    ProcedureAndQuery,
+    QueryByProcedure,
+    procedure_lineage_processor,
+)
 from metadata.utils.filters import (
     filter_by_database,
     filter_by_schema,
@@ -50,53 +45,10 @@ from metadata.utils.filters import (
 from metadata.utils.helpers import pprint_format_object
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.stored_procedures import get_procedure_name_from_call
-from metadata.utils.time_utils import datetime_to_timestamp
 
 logger = ingestion_logger()
 
 
-class QueryByProcedure(BaseModel):
-    """
-    Query(ies) executed by each stored procedure
-    """
-
-    procedure_name: str = Field(None, alias="PROCEDURE_NAME")
-    query_type: str = Field(..., alias="QUERY_TYPE")
-    query_database_name: Optional[str] = Field(None, alias="QUERY_DATABASE_NAME")
-    query_schema_name: Optional[str] = Field(None, alias="QUERY_SCHEMA_NAME")
-    procedure_text: str = Field(..., alias="PROCEDURE_TEXT")
-    procedure_start_time: datetime = Field(..., alias="PROCEDURE_START_TIME")
-    procedure_end_time: datetime = Field(..., alias="PROCEDURE_END_TIME")
-    query_start_time: Optional[datetime] = Field(None, alias="QUERY_START_TIME")
-    query_duration: Optional[float] = Field(None, alias="QUERY_DURATION")
-    query_text: str = Field(..., alias="QUERY_TEXT")
-    query_user_name: Optional[str] = Field(None, alias="QUERY_USER_NAME")
-
-    model_config = ConfigDict(populate_by_name=True)
-
-
-class ProcedureAndQuery(BaseModel):
-    """
-    Model to hold the procedure and its queries
-    """
-
-    procedure: StoredProcedure
-    query_by_procedure: QueryByProcedure
-
-    model_config = ConfigDict(populate_by_name=True)
-
-
-class ProcedureAndProcedureGraph(BaseModel):
-    """
-    Model to hold the procedure and its graph
-    """
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    procedure: StoredProcedure
-    graph: DiGraph
-
-
 class StoredProcedureLineageMixin(ABC):
     """
     The full flow is:
@@ -118,33 +70,28 @@ class StoredProcedureLineageMixin(ABC):
     metadata: OpenMetadata
 
     @abstractmethod
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+    def get_stored_procedure_sql_statement(self) -> str:
         """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
+        Return the SQL statement to get the stored procedure queries
         """
 
-    def procedure_queries_dict(self, query: str) -> Dict[str, List[QueryByProcedure]]:
+    def yield_stored_procedure_queries(self) -> Iterator[QueryByProcedure]:
         """
-        Cache the queries ran for the stored procedures in the last `queryLogDuration` days.
-
-        We will run this for each different and db name.
-
-        The dictionary key will be the case-insensitive procedure name.
+        Yield query and stored procedure object for lineage processing.
         """
+        query = self.get_stored_procedure_sql_statement()
         results = self.engine.execute(query).all()
-        queries_dict = defaultdict(list)
 
         for row in results:
             try:
                 query_by_procedure = QueryByProcedure.model_validate(dict(row))
-                procedure_name = (
+                query_by_procedure.procedure_name = (
                     query_by_procedure.procedure_name
                     or get_procedure_name_from_call(
                         query_text=query_by_procedure.procedure_text,
                     )
                 )
-                queries_dict[procedure_name].append(query_by_procedure)
+                yield query_by_procedure
             except Exception as exc:
                 self.status.failed(
                     StackTraceError(
@@ -154,147 +101,10 @@ class StoredProcedureLineageMixin(ABC):
                     )
                 )
 
-        queries_count_per_procedure = {
-            procedure_name: len(queries)
-            for procedure_name, queries in queries_dict.items()
-        }
-        logger.info(
-            f"Count of queries executed for stored procedures: {sum(queries_count_per_procedure.values())}"
-        )
-        logger.info(
-            f"Count of queries per stored procedure: {pprint_format_object(queries_count_per_procedure)}"
-        )
-
-        return queries_dict
-
-    @staticmethod
-    def is_lineage_query(query_type: str, query_text: str) -> bool:
-        """Check if it's worth it to parse the query for lineage"""
-
-        logger.debug(
-            f"Validating query lineage for type [{query_type}] and text [{query_text}]"
-        )
-
-        if query_type in ("MERGE", "UPDATE", "CREATE_TABLE_AS_SELECT"):
-            return True
-
-        if query_type == "INSERT" and re.search(
-            "^.*insert.*into.*select.*$", query_text.replace("\n", " "), re.IGNORECASE
-        ):
-            return True
-
-        return False
-
-    def _yield_procedure_lineage(
-        self, query_by_procedure: QueryByProcedure, procedure: StoredProcedure
-    ) -> Iterable[Either[AddLineageRequest]]:
-        """Add procedure lineage from its query"""
-        self.stored_procedure_query_lineage = False
-        if self.is_lineage_query(
-            query_type=query_by_procedure.query_type,
-            query_text=query_by_procedure.query_text,
-        ):
-            graph = None
-            if self.source_config.enableTempTableLineage:
-                if not self.procedure_graph_map.get(procedure.fullyQualifiedName.root):
-                    # Map to store the directed graph for each procedure with its FQN as key
-                    self.procedure_graph_map[
-                        procedure.fullyQualifiedName.root
-                    ] = ProcedureAndProcedureGraph(procedure=procedure, graph=DiGraph())
-
-                graph = self.procedure_graph_map.get(
-                    procedure.fullyQualifiedName.root
-                ).graph
-
-            self.stored_procedure_query_lineage = True
-            for either_lineage in get_lineage_by_query(
-                self.metadata,
-                query=query_by_procedure.query_text,
-                service_name=self.service_name,
-                database_name=query_by_procedure.query_database_name,
-                schema_name=query_by_procedure.query_schema_name,
-                dialect=ConnectionTypeDialectMapper.dialect_of(
-                    self.service_connection.type.value
-                ),
-                timeout_seconds=self.source_config.parsingTimeoutLimit,
-                lineage_source=LineageSource.QueryLineage,
-                graph=graph,
-            ):
-                if (
-                    either_lineage.left is None
-                    and either_lineage.right.edge.lineageDetails
-                ):
-                    either_lineage.right.edge.lineageDetails.pipeline = EntityReference(
-                        id=procedure.id,
-                        type="storedProcedure",
-                    )
-
-                yield either_lineage
-
-    def yield_procedure_query(
-        self, query_by_procedure: QueryByProcedure, procedure: StoredProcedure
-    ) -> Iterable[Either[CreateQueryRequest]]:
-        """Check the queries triggered by the procedure and add their lineage, if any"""
-
-        yield Either(
-            right=CreateQueryRequest(
-                query=SqlQuery(query_by_procedure.query_text),
-                query_type=query_by_procedure.query_type,
-                duration=query_by_procedure.query_duration,
-                queryDate=Timestamp(
-                    root=datetime_to_timestamp(
-                        query_by_procedure.query_start_time, True
-                    )
-                ),
-                triggeredBy=EntityReference(
-                    id=procedure.id,
-                    type="storedProcedure",
-                ),
-                processedLineage=bool(self.stored_procedure_query_lineage),
-                service=self.service_name,
-            )
-        )
-
-    def procedure_lineage_processor(
-        self, procedure_and_queries: List[ProcedureAndQuery], queue: Queue
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        for procedure_and_query in procedure_and_queries:
-            try:
-                for lineage in self._yield_procedure_lineage(
-                    query_by_procedure=procedure_and_query.query_by_procedure,
-                    procedure=procedure_and_query.procedure,
-                ):
-                    if lineage and lineage.right is not None:
-                        queue.put(
-                            Either(
-                                right=OMetaLineageRequest(
-                                    override_lineage=False,
-                                    lineage_request=lineage.right,
-                                    entity=StoredProcedure,
-                                    entity_fqn=procedure_and_query.procedure.fullyQualifiedName.root,
-                                )
-                            )
-                        )
-                    else:
-                        queue.put(lineage)
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Could not get lineage for store procedure '{procedure_and_query.procedure.fullyQualifiedName}' due to [{exc}]."
-                )
-            try:
-                for lineage in self.yield_procedure_query(
-                    query_by_procedure=procedure_and_query.query_by_procedure,
-                    procedure=procedure_and_query.procedure,
-                ):
-                    queue.put(lineage)
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Could not get query for store procedure '{procedure_and_query.procedure.fullyQualifiedName}' due to [{exc}]."
-                )
-
-    def procedure_lineage_generator(self) -> Iterable[ProcedureAndQuery]:
+    def procedure_lineage_producer(self) -> Iterator[ProcedureAndQuery]:
+        """
+        Generate lineage for a list of stored procedures
+        """
         query = {
             "query": {
                 "bool": {
@@ -321,9 +131,12 @@ class StoredProcedureLineageMixin(ABC):
             )
         query_filter = json.dumps(query)
         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
+
+        procedures_dict = {}
+        queries = self.yield_stored_procedure_queries()
+        queries_count_per_procedure = defaultdict(int)
+
+        # Get the filtered list of stored procedure to process
         for procedure in (
             self.metadata.paginate_es(
                 entity=StoredProcedure, query_filter=query_filter, size=10
@@ -351,20 +164,52 @@ class StoredProcedureLineageMixin(ABC):
                     )
                     continue
                 logger.debug(f"Processing Lineage for [{procedure.name}]")
-                for query_by_procedure in (
-                    queries_dict.get(procedure.name.root.lower()) or []
-                ):
-                    yield ProcedureAndQuery(
-                        procedure=procedure, query_by_procedure=query_by_procedure
-                    )
+                procedures_dict[procedure.name.root.lower()] = procedure
+
+        # Yield the ProcedureAndQuery for filtered stored procedure
+        for query_by_procedure in queries:
+            if not query_by_procedure.procedure_name:
+                continue
+
+            procedure_name = query_by_procedure.procedure_name.lower()
+            queries_count_per_procedure[procedure_name] += 1
+
+            if procedure_name in procedures_dict:
+                yield ProcedureAndQuery(
+                    procedure=procedures_dict[procedure_name],
+                    query_by_procedure=query_by_procedure,
+                )
+
+        logger.info(
+            f"Count of queries executed for stored procedures: {sum(queries_count_per_procedure.values())}"
+        )
+        logger.info(
+            f"Count of queries per stored procedure: {pprint_format_object(dict(queries_count_per_procedure))}"
+        )
 
     def yield_procedure_lineage(
         self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+    ) -> Iterator[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
         """Get all the queries and procedures list and yield them"""
         logger.info("Processing Lineage for Stored Procedures")
-        producer_fn = self.procedure_lineage_generator
-        processor_fn = self.procedure_lineage_processor
-        yield from self.generate_lineage_in_thread(
-            producer_fn, processor_fn, max_threads=self.source_config.threads
+        producer_fn = self.procedure_lineage_producer
+        processor_fn = procedure_lineage_processor
+        dialect = ConnectionTypeDialectMapper.dialect_of(
+            self.service_connection.type.value
+        )
+        args = (
+            self.metadata,
+            self.service_name,
+            dialect,
+            self.source_config.processCrossDatabaseLineage,
+            self.source_config.crossDatabaseServiceNames,
+            self.source_config.parsingTimeoutLimit,
+            self.procedure_graph_map,
+            self.source_config.enableTempTableLineage,
+        )
+        yield from self.generate_lineage_with_processes(
+            producer_fn,
+            processor_fn,
+            args,
+            max_threads=self.source_config.threads,
         )

metadata/ingestion/source/database/timescale/__init__.py

@@ -0,0 +1,3 @@
+from metadata.ingestion.source.database.timescale.metadata import TimescaleSource
+
+__all__ = ["TimescaleSource"]

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

@@ -0,0 +1,103 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Source connection handler for TimescaleDB
+"""
+from typing import Optional
+
+from sqlalchemy.engine import Engine
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.timescaleConnection import (
+    TimescaleConnection as TimescaleConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    get_connection_url_common,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.connections.test_connections import test_connection_db_common
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.postgres.queries import (
+    POSTGRES_GET_DATABASE,
+    POSTGRES_TEST_GET_QUERIES,
+    POSTGRES_TEST_GET_TAGS,
+)
+from metadata.ingestion.source.database.postgres.utils import (
+    get_postgres_time_column_name,
+)
+from metadata.utils.constants import THREE_MIN
+from metadata.utils.credentials import get_azure_access_token
+
+
+class TimescaleConnection(BaseConnection[TimescaleConnectionConfig, Engine]):
+    def _get_client(self) -> Engine:
+        """
+        Return the SQLAlchemy Engine for TimescaleDB.
+        TimescaleDB is a PostgreSQL extension, so we use the same connection logic.
+        """
+        connection = self.service_connection
+
+        if isinstance(connection.authType, AzureConfigurationSource):
+            access_token = get_azure_access_token(connection.authType)
+            connection.authType = BasicAuth(password=access_token)
+        return create_generic_db_connection(
+            connection=connection,
+            get_connection_url_fn=get_connection_url_common,
+            get_connection_args_fn=get_connection_args_common,
+        )
+
+    def get_connection_dict(self) -> dict:
+        """
+        Return the connection dictionary for this service.
+        """
+        raise NotImplementedError(
+            "get_connection_dict is not implemented for TimescaleDB"
+        )
+
+    def test_connection(
+        self,
+        metadata: OpenMetadata,
+        automation_workflow: Optional[AutomationWorkflow] = None,
+        timeout_seconds: Optional[int] = THREE_MIN,
+    ) -> TestConnectionResult:
+        """
+        Test connection. This can be executed either as part
+        of a metadata workflow or during an Automation Workflow
+        """
+        queries = {
+            "GetQueries": POSTGRES_TEST_GET_QUERIES.format(
+                time_column_name=get_postgres_time_column_name(engine=self.client),
+            ),
+            "GetDatabases": POSTGRES_GET_DATABASE,
+            "GetTags": POSTGRES_TEST_GET_TAGS,
+        }
+        return test_connection_db_common(
+            metadata=metadata,
+            engine=self.client,
+            service_connection=self.service_connection,
+            automation_workflow=automation_workflow,
+            timeout_seconds=timeout_seconds,
+            queries=queries,
+        )