acryl-datahub 1.2.0.1rc1__py3-none-any.whl → 1.2.0.2rc2__py3-none-any.whl

datahub/ingestion/source/sql/postgres.py

@@ -36,6 +36,15 @@ from datahub.ingestion.source.sql.sql_common import (
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
+from datahub.ingestion.source.sql.sql_utils import (
+    gen_database_key,
+    gen_schema_key,
+)
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    BaseProcedure,
+    generate_procedure_container_workunits,
+    generate_procedure_workunits,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     ArrayTypeClass,
     BytesTypeClass,
@@ -123,6 +132,15 @@ class PostgresConfig(BasePostgresConfig):
             "Note: this is not used if `database` or `sqlalchemy_uri` are provided."
         ),
     )
+    include_stored_procedures: bool = Field(
+        default=True,
+        description="Include ingest of stored procedures.",
+    )
+    procedure_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for stored procedures to filter in ingestion."
+        "Specify regex to match the entire procedure name in database.schema.procedure_name format. e.g. to match all procedures starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+    )
 
 
 @platform_name("Postgres")
@@ -135,7 +153,7 @@ class PostgresSource(SQLAlchemySource):
     """
     This plugin extracts the following:
 
-    - Metadata for databases, schemas, views, and tables
+    - Metadata for databases, schemas, views, tables, and stored procedures
     - Column types associated with each table
     - Also supports PostGIS extensions
     - Table, row, and column statistics via optional SQL profiling
@@ -291,3 +309,174 @@ class PostgresSource(SQLAlchemySource):
                     ] = row.table_size
         except Exception as e:
             logger.error(f"failed to fetch profile metadata: {e}")
+
+    def get_schema_level_workunits(
+        self,
+        inspector: Inspector,
+        schema: str,
+        database: str,
+    ) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        yield from super().get_schema_level_workunits(
+            inspector=inspector,
+            schema=schema,
+            database=database,
+        )
+
+        if self.config.include_stored_procedures:
+            try:
+                yield from self.loop_stored_procedures(inspector, schema, self.config)
+            except Exception as e:
+                self.report.failure(
+                    title="Failed to list stored procedures for schema",
+                    message="An error occurred while listing procedures for the schema.",
+                    context=f"{database}.{schema}",
+                    exc=e,
+                )
+
+    def loop_stored_procedures(
+        self,
+        inspector: Inspector,
+        schema: str,
+        config: PostgresConfig,
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Loop schema data for get stored procedures as dataJob-s.
+        """
+        db_name = self.get_db_name(inspector)
+
+        procedures = self.fetch_procedures_for_schema(inspector, schema, db_name)
+        if procedures:
+            yield from self._process_procedures(procedures, db_name, schema)
+
+    def fetch_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        try:
+            raw_procedures: List[BaseProcedure] = self.get_procedures_for_schema(
+                inspector, schema, db_name
+            )
+            procedures: List[BaseProcedure] = []
+            for procedure in raw_procedures:
+                procedure_qualified_name = self.get_identifier(
+                    schema=schema,
+                    entity=procedure.name,
+                    inspector=inspector,
+                )
+
+                if not self.config.procedure_pattern.allowed(procedure_qualified_name):
+                    self.report.report_dropped(procedure_qualified_name)
+                else:
+                    procedures.append(procedure)
+            return procedures
+        except Exception as e:
+            self.report.warning(
+                title="Failed to get procedures for schema",
+                message="An error occurred while fetching procedures for the schema.",
+                context=f"{db_name}.{schema}",
+                exc=e,
+            )
+            return []
+
+    def get_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        """
+        Get stored procedures for a specific schema.
+        """
+        base_procedures = []
+        with inspector.engine.connect() as conn:
+            procedures = conn.execute(
+                """
+                    SELECT
+                        p.proname AS name,
+                        l.lanname AS language,
+                        pg_get_function_arguments(p.oid) AS arguments,
+                        pg_get_functiondef(p.oid) AS definition,
+                        obj_description(p.oid, 'pg_proc') AS comment
+                    FROM
+                        pg_proc p
+                    JOIN
+                        pg_namespace n ON n.oid = p.pronamespace
+                    JOIN
+                        pg_language l ON l.oid = p.prolang
+                    WHERE
+                        p.prokind = 'p'
+                        AND n.nspname = '"""
+                + schema
+                + """';
+                    """
+            )
+
+            procedure_rows = list(procedures)
+            for row in procedure_rows:
+                base_procedures.append(
+                    BaseProcedure(
+                        name=row.name,
+                        language=row.language,
+                        argument_signature=row.arguments,
+                        return_type=None,
+                        procedure_definition=row.definition,
+                        created=None,
+                        last_altered=None,
+                        comment=row.comment,
+                        extra_properties=None,
+                    )
+                )
+            return base_procedures
+
+    def _process_procedures(
+        self,
+        procedures: List[BaseProcedure],
+        db_name: str,
+        schema: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if procedures:
+            yield from generate_procedure_container_workunits(
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+            )
+        for procedure in procedures:
+            yield from self._process_procedure(procedure, schema, db_name)
+
+    def _process_procedure(
+        self,
+        procedure: BaseProcedure,
+        schema: str,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        try:
+            yield from generate_procedure_workunits(
+                procedure=procedure,
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_resolver=self.get_schema_resolver(),
+            )
+        except Exception as e:
+            self.report.warning(
+                title="Failed to emit stored procedure",
+                message="An error occurred while emitting stored procedure",
+                context=procedure.name,
+                exc=e,
+            )

datahub/ingestion/source/sql_queries.py

@@ -4,7 +4,7 @@ import os
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import partial
-from typing import Iterable, List, Optional, Set
+from typing import Iterable, List, Optional, Union
 
 from pydantic import Field
 
@@ -14,9 +14,8 @@ from datahub.configuration.source_common import (
 )
 from datahub.emitter.mce_builder import (
     make_dataset_urn_with_platform_instance,
-    make_user_urn,
 )
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -35,8 +34,15 @@ from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
+from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
+from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
-from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
+from datahub.sql_parsing.sql_parsing_aggregator import (
+    KnownQueryLineageInfo,
+    ObservedQuery,
+    SqlAggregatorReport,
+    SqlParsingAggregator,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -67,28 +73,19 @@ class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
         default=None,
     )
     override_dialect: Optional[str] = Field(
-        description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
+        description="DEPRECATED: This field is ignored. SQL dialect detection is now handled automatically by the SQL parsing aggregator based on the platform.",
         default=None,
+        hidden_from_docs=True,
     )
 
 
-class SqlQueriesSourceReport(SourceReport):
-    num_queries_parsed: int = 0
-    num_table_parse_failures: int = 0
-    num_column_parse_failures: int = 0
+@dataclass
+class SqlQueriesSourceReport(SourceReport, IngestionStageReport):
+    num_entries_processed: int = 0
+    num_entries_failed: int = 0
+    num_queries_aggregator_failures: int = 0
 
-    def compute_stats(self) -> None:
-        super().compute_stats()
-        self.table_failure_rate = (
-            f"{self.num_table_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
-        self.column_failure_rate = (
-            f"{self.num_column_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
+    sql_aggregator: Optional[SqlAggregatorReport] = None
 
 
 @platform_name("SQL Queries")
@@ -107,15 +104,15 @@ class SqlQueriesSource(Source):
     - user (optional): string - The user who ran the query.
     This user value will be directly converted into a DataHub user urn.
     - operation_type (optional): string - Platform-specific operation type, used if the operation type can't be parsed.
+    - session_id (optional): string - Session identifier for temporary table resolution across queries.
     - downstream_tables (optional): string[] - Fallback list of tables that the query writes to,
      used if the query can't be parsed.
     - upstream_tables (optional): string[] - Fallback list of tables the query reads from,
      used if the query can't be parsed.
     """
 
-    urns: Optional[Set[str]]
-    schema_resolver: SchemaResolver
-    builder: SqlParsingBuilder
+    schema_resolver: Optional[SchemaResolver]
+    aggregator: SqlParsingAggregator
 
     def __init__(self, ctx: PipelineContext, config: SqlQueriesSourceConfig):
         if not ctx.graph:
@@ -128,22 +125,36 @@ class SqlQueriesSource(Source):
         self.config = config
         self.report = SqlQueriesSourceReport()
 
-        self.builder = SqlParsingBuilder(usage_config=self.config.usage)
-
         if self.config.use_schema_resolver:
+            # TODO: `initialize_schema_resolver_from_datahub` does a  bulk initialization by fetching all schemas
+            # for the given platform, platform instance, and env. Instead this should be configurable:
+            # bulk initialization vs lazy on-demand schema fetching.
             self.schema_resolver = self.graph.initialize_schema_resolver_from_datahub(
                 platform=self.config.platform,
                 platform_instance=self.config.platform_instance,
                 env=self.config.env,
             )
-            self.urns = self.schema_resolver.get_urns()
         else:
-            self.schema_resolver = self.graph._make_schema_resolver(
-                platform=self.config.platform,
-                platform_instance=self.config.platform_instance,
-                env=self.config.env,
-            )
-            self.urns = None
+            self.schema_resolver = None
+
+        self.aggregator = SqlParsingAggregator(
+            platform=self.config.platform,
+            platform_instance=self.config.platform_instance,
+            env=self.config.env,
+            schema_resolver=self.schema_resolver,
+            eager_graph_load=False,
+            generate_lineage=True,  # TODO: make this configurable
+            generate_queries=True,  # TODO: make this configurable
+            generate_query_subject_fields=True,  # TODO: make this configurable
+            generate_query_usage_statistics=True,  # This enables publishing SELECT query entities, otherwise only mutation queries are published
+            generate_usage_statistics=True,
+            generate_operations=True,  # TODO: make this configurable
+            usage_config=self.config.usage,
+            is_temp_table=None,
+            is_allowed_table=None,
+            format_queries=False,
+        )
+        self.report.sql_aggregator = self.aggregator.report
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "SqlQueriesSource":
@@ -156,68 +167,91 @@ class SqlQueriesSource(Source):
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [partial(auto_workunit_reporter, self.get_report())]
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         logger.info(f"Parsing queries from {os.path.basename(self.config.query_file)}")
+
+        with self.report.new_stage("Collecting queries from file"):
+            queries = list(self._parse_query_file())
+            logger.info(f"Collected {len(queries)} queries for processing")
+
+        with self.report.new_stage("Processing queries through SQL parsing aggregator"):
+            for query_entry in queries:
+                self._add_query_to_aggregator(query_entry)
+
+        with self.report.new_stage("Generating metadata work units"):
+            logger.info("Generating workunits from SQL parsing aggregator")
+            yield from self.aggregator.gen_metadata()
+
+    def _parse_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse the query file and yield QueryEntry objects."""
         with open(self.config.query_file) as f:
             for line in f:
                 try:
                     query_dict = json.loads(line, strict=False)
                     entry = QueryEntry.create(query_dict, config=self.config)
-                    yield from self._process_query(entry)
+                    self.report.num_entries_processed += 1
+                    if self.report.num_entries_processed % 1000 == 0:
+                        logger.info(
+                            f"Processed {self.report.num_entries_processed} query entries"
+                        )
+                    yield entry
                 except Exception as e:
-                    logger.warning("Error processing query", exc_info=True)
-                    self.report.report_warning("process-query", str(e))
-
-        logger.info("Generating workunits")
-        yield from self.builder.gen_workunits()
+                    self.report.num_entries_failed += 1
+                    self.report.warning(
+                        title="Error processing query",
+                        message="Query skipped due to parsing error",
+                        context=line.strip(),
+                        exc=e,
+                    )
 
-    def _process_query(self, entry: "QueryEntry") -> Iterable[MetadataWorkUnit]:
-        self.report.num_queries_parsed += 1
-        if self.report.num_queries_parsed % 1000 == 0:
-            logger.info(f"Parsed {self.report.num_queries_parsed} queries")
-
-        result = sqlglot_lineage(
-            sql=entry.query,
-            schema_resolver=self.schema_resolver,
-            default_db=self.config.default_db,
-            default_schema=self.config.default_schema,
-            override_dialect=self.config.override_dialect,
-        )
-        if result.debug_info.table_error:
-            logger.info(f"Error parsing table lineage, {result.debug_info.table_error}")
-            self.report.num_table_parse_failures += 1
-            for downstream_urn in set(entry.downstream_tables):
-                self.builder.add_lineage(
-                    downstream_urn=downstream_urn,
-                    upstream_urns=entry.upstream_tables,
-                    timestamp=entry.timestamp,
-                    user=entry.user,
+    def _add_query_to_aggregator(self, query_entry: "QueryEntry") -> None:
+        """Add a query to the SQL parsing aggregator."""
+        try:
+            # If we have explicit lineage, use it directly
+            if query_entry.upstream_tables or query_entry.downstream_tables:
+                logger.debug("Using explicit lineage from query file")
+                for downstream_table in query_entry.downstream_tables:
+                    known_lineage = KnownQueryLineageInfo(
+                        query_text=query_entry.query,
+                        downstream=downstream_table,
+                        upstreams=query_entry.upstream_tables,
+                        timestamp=query_entry.timestamp,
+                        session_id=query_entry.session_id,
+                    )
+                    self.aggregator.add_known_query_lineage(known_lineage)
+            else:
+                # No explicit lineage, rely on parsing
+                observed_query = ObservedQuery(
+                    query=query_entry.query,
+                    timestamp=query_entry.timestamp,
+                    user=query_entry.user,
+                    session_id=query_entry.session_id,
+                    default_db=self.config.default_db,
+                    default_schema=self.config.default_schema,
                 )
-            return
-        elif result.debug_info.column_error:
-            logger.debug(
-                f"Error parsing column lineage, {result.debug_info.column_error}"
+                self.aggregator.add_observed_query(observed_query)
+
+        except Exception as e:
+            self.report.num_queries_aggregator_failures += 1
+            self.report.warning(
+                title="Error adding query to aggregator",
+                message="Query skipped due to failure when adding query to SQL parsing aggregator",
+                context=query_entry.query,
+                exc=e,
             )
-            self.report.num_column_parse_failures += 1
-
-        yield from self.builder.process_sql_parsing_result(
-            result,
-            query=entry.query,
-            query_timestamp=entry.timestamp,
-            user=entry.user,
-            custom_operation_type=entry.operation_type,
-            include_urns=self.urns,
-        )
 
 
 @dataclass
 class QueryEntry:
     query: str
     timestamp: Optional[datetime]
-    user: Optional[str]
+    user: Optional[CorpUserUrn]
     operation_type: Optional[str]
     downstream_tables: List[str]
     upstream_tables: List[str]
+    session_id: Optional[str] = None
 
     @classmethod
     def create(
@@ -230,7 +264,7 @@ class QueryEntry:
                 if "timestamp" in entry_dict
                 else None
             ),
-            user=make_user_urn(entry_dict["user"]) if "user" in entry_dict else None,
+            user=CorpUserUrn(entry_dict["user"]) if "user" in entry_dict else None,
             operation_type=entry_dict.get("operation_type"),
             downstream_tables=[
                 make_dataset_urn_with_platform_instance(
@@ -250,4 +284,5 @@ class QueryEntry:
                 )
                 for table in entry_dict.get("upstream_tables", [])
             ],
+            session_id=entry_dict.get("session_id"),
         )

datahub/ingestion/source/unity/proxy.py

@@ -521,9 +521,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_schema_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching schema tags for catalog: {catalog}")
+        logger.info(f"Fetching schema tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.schema_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.schema_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -544,9 +544,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_catalog_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching table tags for catalog: {catalog}")
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.catalog_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.catalog_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -566,9 +566,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_table_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching table tags for catalog: {catalog}")
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.table_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.table_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -589,9 +589,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_column_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching column tags for catalog: {catalog}")
+        logger.info(f"Fetching column tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.column_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.column_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}