acryl-datahub 1.0.0rc17__py3-none-any.whl → 1.0.0.1__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_query.py

@@ -164,6 +164,23 @@ class SnowflakeQuery:
         and table_type in ('BASE TABLE', 'EXTERNAL TABLE')
         order by table_schema, table_name"""
 
+    @staticmethod
+    def procedures_for_database(db_name: Optional[str]) -> str:
+        db_clause = f'"{db_name}".' if db_name is not None else ""
+        return f"""
+        SELECT procedure_catalog AS "PROCEDURE_CATALOG",
+        procedure_schema AS "PROCEDURE_SCHEMA",
+        procedure_name AS "PROCEDURE_NAME",
+        procedure_language AS "PROCEDURE_LANGUAGE",
+        argument_signature AS "ARGUMENT_SIGNATURE",
+        data_type AS "PROCEDURE_RETURN_TYPE",
+        procedure_definition AS "PROCEDURE_DEFINITION",
+        created AS "CREATED",
+        last_altered AS "LAST_ALTERED",
+        comment AS "COMMENT"
+        FROM {db_clause}information_schema.procedures
+        order by procedure_schema, procedure_name"""
+
     @staticmethod
     def get_all_tags():
         return """

datahub/ingestion/source/snowflake/snowflake_report.py

@@ -105,6 +105,7 @@ class SnowflakeV2Report(
     databases_scanned: int = 0
     tags_scanned: int = 0
     streams_scanned: int = 0
+    procedures_scanned: int = 0
 
     include_usage_stats: bool = False
     include_operational_stats: bool = False
@@ -163,6 +164,8 @@ class SnowflakeV2Report(
             self.tags_scanned += 1
         elif ent_type == "stream":
             self.streams_scanned += 1
+        elif ent_type == "procedure":
+            self.procedures_scanned += 1
         else:
             raise KeyError(f"Unknown entity {ent_type}.")
 

datahub/ingestion/source/snowflake/snowflake_schema.py

@@ -14,6 +14,7 @@ from datahub.ingestion.source.snowflake.snowflake_query import (
     SnowflakeQuery,
 )
 from datahub.ingestion.source.sql.sql_generic import BaseColumn, BaseTable, BaseView
+from datahub.ingestion.source.sql.stored_procedures.base import BaseProcedure
 from datahub.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.prefix_batch_builder import PrefixGroup, build_prefix_batches
 from datahub.utilities.serialized_lru_cache import serialized_lru_cache
@@ -714,3 +715,31 @@ class SnowflakeDataDictionary(SupportsAsObj):
                 stream_pagination_marker = stream_name
 
         return streams
+
+    @serialized_lru_cache(maxsize=1)
+    def get_procedures_for_database(
+        self, db_name: str
+    ) -> Dict[str, List[BaseProcedure]]:
+        procedures: Dict[str, List[BaseProcedure]] = {}
+        cur = self.connection.query(
+            SnowflakeQuery.procedures_for_database(db_name),
+        )
+
+        for procedure in cur:
+            if procedure["PROCEDURE_SCHEMA"] not in procedures:
+                procedures[procedure["PROCEDURE_SCHEMA"]] = []
+
+            procedures[procedure["PROCEDURE_SCHEMA"]].append(
+                BaseProcedure(
+                    name=procedure["PROCEDURE_NAME"],
+                    language=procedure["PROCEDURE_LANGUAGE"],
+                    argument_signature=procedure["ARGUMENT_SIGNATURE"],
+                    return_type=procedure["PROCEDURE_RETURN_TYPE"],
+                    procedure_definition=procedure["PROCEDURE_DEFINITION"],
+                    created=procedure["CREATED"],
+                    last_altered=procedure["LAST_ALTERED"],
+                    comment=procedure["COMMENT"],
+                    extra_properties=None,
+                )
+            )
+        return procedures

datahub/ingestion/source/snowflake/snowflake_schema_gen.py

@@ -41,6 +41,7 @@ from datahub.ingestion.source.snowflake.snowflake_query import SnowflakeQuery
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
 from datahub.ingestion.source.snowflake.snowflake_schema import (
     SCHEMA_PARALLELISM,
+    BaseProcedure,
     SnowflakeColumn,
     SnowflakeDatabase,
     SnowflakeDataDictionary,
@@ -63,12 +64,14 @@ from datahub.ingestion.source.snowflake.snowflake_utils import (
 from datahub.ingestion.source.sql.sql_utils import (
     add_table_to_schema_container,
     gen_database_container,
-    gen_database_key,
     gen_schema_container,
-    gen_schema_key,
     get_dataplatform_instance_aspect,
     get_domain_wu,
 )
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    generate_procedure_container_workunits,
+    generate_procedure_workunits,
+)
 from datahub.ingestion.source_report.ingestion_stage import (
     EXTERNAL_TABLE_DDL_LINEAGE,
     LINEAGE_EXTRACTION,
@@ -448,10 +451,15 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         if self.config.include_streams:
             self.report.num_get_streams_for_schema_queries += 1
             streams = self.fetch_streams_for_schema(
-                snowflake_schema, db_name, schema_name
+                snowflake_schema,
+                db_name,
             )
             yield from self._process_streams(streams, snowflake_schema, db_name)
 
+        if self.config.include_procedures:
+            procedures = self.fetch_procedures_for_schema(snowflake_schema, db_name)
+            yield from self._process_procedures(procedures, snowflake_schema, db_name)
+
         if self.config.include_technical_schema and snowflake_schema.tags:
             yield from self._process_tags_in_schema(snowflake_schema)
 
@@ -536,6 +544,26 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         for stream in streams:
             yield from self._process_stream(stream, snowflake_schema, db_name)
 
+    def _process_procedures(
+        self,
+        procedures: List[BaseProcedure],
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if self.config.include_technical_schema:
+            if procedures:
+                yield from generate_procedure_container_workunits(
+                    self.identifiers.gen_database_key(
+                        db_name,
+                    ),
+                    self.identifiers.gen_schema_key(
+                        db_name=db_name,
+                        schema_name=snowflake_schema.name,
+                    ),
+                )
+            for procedure in procedures:
+                yield from self._process_procedure(procedure, snowflake_schema, db_name)
+
     def _process_tags_in_schema(
         self, snowflake_schema: SnowflakeSchema
     ) -> Iterable[MetadataWorkUnit]:
@@ -819,13 +847,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             entityUrn=dataset_urn, aspect=dataset_properties
         ).as_workunit()
 
-        schema_container_key = gen_schema_key(
-            db_name=self.snowflake_identifier(db_name),
-            schema=self.snowflake_identifier(schema_name),
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
-        )
+        schema_container_key = self.identifiers.gen_schema_key(db_name, schema_name)
 
         if self.config.extract_tags_as_structured_properties:
             yield from self.gen_column_tags_as_structured_properties(dataset_urn, table)
@@ -1094,11 +1116,8 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
     def gen_database_containers(
         self, database: SnowflakeDatabase
     ) -> Iterable[MetadataWorkUnit]:
-        database_container_key = gen_database_key(
-            self.snowflake_identifier(database.name),
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
+        database_container_key = self.identifiers.gen_database_key(
+            database.name,
         )
 
         yield from gen_database_container(
@@ -1147,21 +1166,9 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
     def gen_schema_containers(
         self, schema: SnowflakeSchema, db_name: str
     ) -> Iterable[MetadataWorkUnit]:
-        schema_name = self.snowflake_identifier(schema.name)
-        database_container_key = gen_database_key(
-            database=self.snowflake_identifier(db_name),
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
-        )
+        database_container_key = self.identifiers.gen_database_key(db_name)
 
-        schema_container_key = gen_schema_key(
-            db_name=self.snowflake_identifier(db_name),
-            schema=schema_name,
-            platform=self.platform,
-            platform_instance=self.config.platform_instance,
-            env=self.config.env,
-        )
+        schema_container_key = self.identifiers.gen_schema_key(db_name, schema.name)
 
         yield from gen_schema_container(
             name=schema.name,
@@ -1290,13 +1297,13 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             )
 
     def fetch_streams_for_schema(
-        self, snowflake_schema: SnowflakeSchema, db_name: str, schema_name: str
+        self, snowflake_schema: SnowflakeSchema, db_name: str
     ) -> List[SnowflakeStream]:
         try:
             streams: List[SnowflakeStream] = []
-            for stream in self.get_streams_for_schema(schema_name, db_name):
+            for stream in self.get_streams_for_schema(snowflake_schema.name, db_name):
                 stream_identifier = self.identifiers.get_dataset_identifier(
-                    stream.name, schema_name, db_name
+                    stream.name, snowflake_schema.name, db_name
                 )
 
                 self.report.report_entity_scanned(stream_identifier, "stream")
@@ -1310,16 +1317,15 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             snowflake_schema.streams = [stream.name for stream in streams]
             return streams
         except Exception as e:
-            if isinstance(e, SnowflakePermissionError):
-                error_msg = f"Failed to get streams for schema {db_name}.{schema_name}. Please check permissions."
-                raise SnowflakePermissionError(error_msg) from e.__cause__
-            else:
-                self.structured_reporter.warning(
-                    "Failed to get streams for schema",
-                    f"{db_name}.{schema_name}",
-                    exc=e,
-                )
-                return []
+            self.structured_reporter.warning(
+                title="Failed to get streams for schema",
+                message="Please check permissions"
+                if isinstance(e, SnowflakePermissionError)
+                else "",
+                context=f"{db_name}.{snowflake_schema.name}",
+                exc=e,
+            )
+            return []
 
     def get_streams_for_schema(
         self, schema_name: str, db_name: str
@@ -1328,6 +1334,42 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
         return streams.get(schema_name, [])
 
+    def fetch_procedures_for_schema(
+        self, snowflake_schema: SnowflakeSchema, db_name: str
+    ) -> List[BaseProcedure]:
+        try:
+            procedures: List[BaseProcedure] = []
+            for procedure in self.get_procedures_for_schema(snowflake_schema, db_name):
+                procedure_qualified_name = self.identifiers.get_dataset_identifier(
+                    procedure.name, snowflake_schema.name, db_name
+                )
+                self.report.report_entity_scanned(procedure_qualified_name, "procedure")
+
+                if self.filters.is_procedure_allowed(procedure_qualified_name):
+                    procedures.append(procedure)
+                else:
+                    self.report.report_dropped(procedure_qualified_name)
+            return procedures
+        except Exception as e:
+            self.structured_reporter.warning(
+                title="Failed to get procedures for schema",
+                message="Please check permissions"
+                if isinstance(e, SnowflakePermissionError)
+                else "",
+                context=f"{db_name}.{snowflake_schema.name}",
+                exc=e,
+            )
+            return []
+
+    def get_procedures_for_schema(
+        self,
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> List[BaseProcedure]:
+        procedures = self.data_dictionary.get_procedures_for_database(db_name)
+
+        return procedures.get(snowflake_schema.name, [])
+
     def _process_stream(
         self,
         stream: SnowflakeStream,
@@ -1350,6 +1392,34 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
                 "Failed to get columns for stream:", stream.name, exc=e
             )
 
+    def _process_procedure(
+        self,
+        procedure: BaseProcedure,
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        try:
+            # TODO: For CLL, we should process procedures after all tables are processed
+            yield from generate_procedure_workunits(
+                procedure,
+                database_key=self.identifiers.gen_database_key(
+                    db_name,
+                ),
+                schema_key=self.identifiers.gen_schema_key(
+                    db_name, snowflake_schema.name
+                ),
+                schema_resolver=(
+                    self.aggregator._schema_resolver if self.aggregator else None
+                ),
+            )
+        except Exception as e:
+            self.structured_reporter.warning(
+                title="Failed to ingest stored procedure",
+                message="",
+                context=procedure.name,
+                exc=e,
+            )
+
     def get_columns_for_stream(
         self,
         source_object: str,  # Qualified name of source table/view

datahub/ingestion/source/snowflake/snowflake_utils.py

@@ -3,7 +3,10 @@ from functools import cached_property
 from typing import ClassVar, List, Literal, Optional, Tuple
 
 from datahub.configuration.pattern_utils import is_schema_allowed
-from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
+from datahub.emitter.mce_builder import (
+    make_dataset_urn_with_platform_instance,
+)
+from datahub.emitter.mcp_builder import DatabaseKey, SchemaKey
 from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.source.snowflake.constants import (
     SNOWFLAKE_REGION_CLOUD_REGION_MAPPING,
@@ -16,6 +19,7 @@ from datahub.ingestion.source.snowflake.snowflake_config import (
     SnowflakeV2Config,
 )
 from datahub.ingestion.source.snowflake.snowflake_report import SnowflakeV2Report
+from datahub.ingestion.source.sql.sql_utils import gen_database_key, gen_schema_key
 
 
 class SnowflakeStructuredReportMixin(abc.ABC):
@@ -180,6 +184,9 @@ class SnowflakeFilter:
 
         return True
 
+    def is_procedure_allowed(self, procedure_name: str) -> bool:
+        return self.filter_config.procedure_pattern.allowed(procedure_name)
+
 
 def _combine_identifier_parts(
     *, table_name: str, schema_name: str, db_name: str
@@ -330,6 +337,23 @@ class SnowflakeIdentifierBuilder:
             else user_name
         )
 
+    def gen_schema_key(self, db_name: str, schema_name: str) -> SchemaKey:
+        return gen_schema_key(
+            db_name=self.snowflake_identifier(db_name),
+            schema=self.snowflake_identifier(schema_name),
+            platform=self.platform,
+            platform_instance=self.identifier_config.platform_instance,
+            env=self.identifier_config.env,
+        )
+
+    def gen_database_key(self, db_name: str) -> DatabaseKey:
+        return gen_database_key(
+            database=self.snowflake_identifier(db_name),
+            platform=self.platform,
+            platform_instance=self.identifier_config.platform_instance,
+            env=self.identifier_config.env,
+        )
+
 
 class SnowflakeCommonMixin(SnowflakeStructuredReportMixin):
     platform = "snowflake"

datahub/ingestion/source/sql/mssql/job_models.py

@@ -15,6 +15,7 @@ from datahub.ingestion.source.common.subtypes import (
     FlowContainerSubTypes,
     JobContainerSubTypes,
 )
+from datahub.ingestion.source.sql.stored_procedures.base import BaseProcedure
 from datahub.metadata.schema_classes import (
     ContainerClass,
     DataFlowInfoClass,
@@ -135,6 +136,19 @@ class StoredProcedure:
     def escape_full_name(self) -> str:
         return f"[{self.db}].[{self.schema}].[{self.formatted_name}]"
 
+    def to_base_procedure(self) -> BaseProcedure:
+        return BaseProcedure(
+            name=self.formatted_name,
+            procedure_definition=self.code,
+            created=None,
+            last_altered=None,
+            comment=None,
+            argument_signature=None,
+            return_type=None,
+            language="SQL",
+            extra_properties=None,
+        )
+
 
 @dataclass
 class JobStep:
@@ -222,7 +236,7 @@ class MSSQLDataJob:
         type = (
             JobContainerSubTypes.MSSQL_JOBSTEP
             if isinstance(self.entity, JobStep)
-            else JobContainerSubTypes.MSSQL_STORED_PROCEDURE
+            else JobContainerSubTypes.STORED_PROCEDURE
         )
         return SubTypesClass(
             typeNames=[type],

datahub/ingestion/source/sql/mssql/source.py

@@ -37,9 +37,6 @@ from datahub.ingestion.source.sql.mssql.job_models import (
     ProcedureParameter,
     StoredProcedure,
 )
-from datahub.ingestion.source.sql.mssql.stored_procedure_lineage import (
-    generate_procedure_lineage,
-)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -50,6 +47,9 @@ from datahub.ingestion.source.sql.sql_config import (
     make_sqlalchemy_uri,
 )
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    generate_procedure_lineage,
+)
 from datahub.utilities.file_backed_collections import FileBackedList
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -65,6 +65,8 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
     # defaults
     host_port: str = Field(default="localhost:1433", description="MSSQL host URL.")
     scheme: str = Field(default="mssql+pytds", description="", hidden_from_docs=True)
+
+    # TODO: rename to include_procedures ?
     include_stored_procedures: bool = Field(
         default=True,
         description="Include ingest of stored procedures. Requires access to the 'sys' schema.",
@@ -763,9 +765,11 @@ class SQLServerSource(SQLAlchemySource):
                 yield from auto_workunit(
                     generate_procedure_lineage(
                         schema_resolver=self.get_schema_resolver(),
-                        procedure=procedure,
+                        procedure=procedure.to_base_procedure(),
                         procedure_job_urn=MSSQLDataJob(entity=procedure).urn,
                         is_temp_table=self.is_temp_table,
+                        default_db=procedure.db,
+                        default_schema=procedure.schema,
                     )
                 )
 

datahub/ingestion/source/sql/oracle.py

@@ -31,7 +31,9 @@ from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     make_sqlalchemy_type,
 )
-from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
+from datahub.ingestion.source.sql.sql_config import (
+    BasicSQLAlchemyConfig,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -71,10 +73,12 @@ class OracleConfig(BasicSQLAlchemyConfig):
         description="Will be set automatically to default value.",
     )
     service_name: Optional[str] = Field(
-        default=None, description="Oracle service name. If using, omit `database`."
+        default=None,
+        description="Oracle service name. If using, omit `database`.",
     )
     database: Optional[str] = Field(
-        default=None, description="If using, omit `service_name`."
+        default=None,
+        description="If using, omit `service_name`.",
     )
     add_database_name_to_urn: Optional[bool] = Field(
         default=False,
@@ -631,7 +635,6 @@ class OracleSource(SQLAlchemySource):
     - Table, row, and column statistics via optional SQL profiling
 
     Using the Oracle source requires that you've also installed the correct drivers; see the [cx_Oracle docs](https://cx-oracle.readthedocs.io/en/latest/user_guide/installation.html). The easiest one is the [Oracle Instant Client](https://www.oracle.com/database/technologies/instant-client.html).
-
     """
 
     config: OracleConfig
@@ -661,6 +664,8 @@ class OracleSource(SQLAlchemySource):
         database name from Connection URL, which does not work when using
         service instead of database.
         In that case, it tries to retrieve the database name by sending a query to the DB.
+
+        Note: This is used as a fallback if database is not specified in the config.
         """
 
         # call default implementation first
@@ -687,7 +692,49 @@ class OracleSource(SQLAlchemySource):
                 # To silent the mypy lint error
                 yield cast(Inspector, inspector)
 
+    def get_db_schema(self, dataset_identifier: str) -> Tuple[Optional[str], str]:
+        """
+        Override the get_db_schema method to ensure proper schema name extraction.
+        This method is used during view lineage extraction to determine the default schema
+        for unqualified table names in view definitions.
+        """
+        try:
+            # Try to get the schema from the dataset identifier
+            parts = dataset_identifier.split(".")
+
+            # Handle the identifier format differently based on add_database_name_to_urn flag
+            if self.config.add_database_name_to_urn:
+                if len(parts) >= 3:
+                    # Format is: database.schema.view when add_database_name_to_urn=True
+                    db_name = parts[-3]
+                    schema_name = parts[-2]
+                    return db_name, schema_name
+                elif len(parts) >= 2:
+                    # Handle the case where database might be missing even with flag enabled
+                    # If we have a database in the config, use that
+                    db_name = str(self.config.database)
+                    schema_name = parts[-2]
+                    return db_name, schema_name
+            else:
+                # Format is: schema.view when add_database_name_to_urn=False
+                if len(parts) >= 2:
+                    # When add_database_name_to_urn is False, don't include database in the result
+                    db_name = None
+                    schema_name = parts[-2]
+                    return db_name, schema_name
+        except Exception as e:
+            logger.warning(
+                f"Error extracting schema from identifier {dataset_identifier}: {e}"
+            )
+
+        # Fall back to parent implementation if our approach fails
+        db_name, schema_name = super().get_db_schema(dataset_identifier)
+        return db_name, schema_name
+
     def get_workunits(self):
+        """
+        Override get_workunits to patch Oracle dialect for custom types.
+        """
         with patch.dict(
             "sqlalchemy.dialects.oracle.base.OracleDialect.ischema_names",
             {klass.__name__: klass for klass in extra_oracle_types},