acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/sql/sqlalchemy_uri.py

@@ -0,0 +1,68 @@
+from typing import Any, Dict, Optional, Tuple
+
+from sqlalchemy.engine import URL
+
+
+def parse_host_port(
+    host_port: str, default_port: Optional[int] = None
+) -> Tuple[str, Optional[int]]:
+    """
+    Parse a host:port string into separate host and port components.
+
+    Args:
+        host_port: String in format "host:port" or just "host"
+        default_port: Optional default port to use if not specified in host_port
+
+    Returns:
+        Tuple of (hostname, port) where port may be None if not specified
+
+    Examples:
+        >>> parse_host_port("localhost:3306")
+        ('localhost', 3306)
+        >>> parse_host_port("localhost")
+        ('localhost', None)
+        >>> parse_host_port("localhost", 5432)
+        ('localhost', 5432)
+        >>> parse_host_port("db.example.com:invalid", 3306)
+        ('db.example.com', 3306)
+    """
+    try:
+        host, port_str = host_port.rsplit(":", 1)
+        port: Optional[int]
+        try:
+            port = int(port_str)
+        except ValueError:
+            # Port is not a valid integer
+            port = default_port
+        return host, port
+    except ValueError:
+        # No colon found, entire string is the hostname
+        return host_port, default_port
+
+
+def make_sqlalchemy_uri(
+    scheme: str,
+    username: Optional[str],
+    password: Optional[str],
+    at: Optional[str],
+    db: Optional[str],
+    uri_opts: Optional[Dict[str, Any]] = None,
+) -> str:
+    host: Optional[str] = None
+    port: Optional[int] = None
+    if at:
+        host, port = parse_host_port(at)
+    if uri_opts:
+        uri_opts = {k: v for k, v in uri_opts.items() if v is not None}
+
+    return str(
+        URL.create(
+            drivername=scheme,
+            username=username,
+            password=password,
+            host=host,
+            port=port,
+            database=db,
+            query=uri_opts or {},
+        )
+    )

datahub/ingestion/source/sql/stored_procedures/base.py

@@ -0,0 +1,253 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Callable, Dict, Iterable, Optional
+
+from datahub.emitter.mce_builder import (
+    DEFAULT_ENV,
+    datahub_guid,
+    make_data_flow_urn,
+    make_data_job_urn,
+    make_data_platform_urn,
+    make_dataplatform_instance_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import DatabaseKey, SchemaKey
+from datahub.ingestion.api.source_helpers import auto_workunit
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import (
+    FlowContainerSubTypes,
+    JobContainerSubTypes,
+)
+from datahub.ingestion.source.sql.stored_procedures.lineage import parse_procedure_code
+from datahub.metadata.schema_classes import (
+    ContainerClass,
+    DataFlowInfoClass,
+    DataJobInfoClass,
+    DataPlatformInstanceClass,
+    DataTransformClass,
+    DataTransformLogicClass,
+    QueryLanguageClass,
+    QueryStatementClass,
+    SubTypesClass,
+)
+from datahub.sql_parsing.schema_resolver import SchemaResolver
+
+
+@dataclass
+class BaseProcedure:
+    name: str
+    procedure_definition: Optional[str]
+    created: Optional[datetime]
+    last_altered: Optional[datetime]
+    comment: Optional[str]
+    argument_signature: Optional[str]
+    return_type: Optional[str]
+    language: str
+    extra_properties: Optional[Dict[str, str]]
+
+    def get_procedure_identifier(
+        self,
+    ) -> str:
+        if self.argument_signature:
+            argument_signature_hash = datahub_guid(
+                dict(argument_signature=self.argument_signature)
+            )
+            return f"{self.name}_{argument_signature_hash}"
+
+        return self.name
+
+    def to_urn(self, database_key: DatabaseKey, schema_key: Optional[SchemaKey]) -> str:
+        return make_data_job_urn(
+            orchestrator=database_key.platform,
+            flow_id=_get_procedure_flow_name(database_key, schema_key),
+            job_id=self.get_procedure_identifier(),
+            cluster=database_key.env or DEFAULT_ENV,
+            platform_instance=database_key.instance,
+        )
+
+
+def _generate_flow_workunits(
+    database_key: DatabaseKey, schema_key: Optional[SchemaKey]
+) -> Iterable[MetadataWorkUnit]:
+    """Generate flow workunits for database and schema"""
+
+    procedure_flow_name = _get_procedure_flow_name(database_key, schema_key)
+
+    flow_urn = make_data_flow_urn(
+        orchestrator=database_key.platform,
+        flow_id=procedure_flow_name,
+        cluster=database_key.env or DEFAULT_ENV,
+        platform_instance=database_key.instance,
+    )
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=flow_urn,
+        aspect=DataFlowInfoClass(
+            name=procedure_flow_name,
+        ),
+    ).as_workunit()
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=flow_urn,
+        aspect=SubTypesClass(
+            typeNames=[FlowContainerSubTypes.MSSQL_PROCEDURE_CONTAINER],
+        ),
+    ).as_workunit()
+
+    if database_key.instance:
+        yield MetadataChangeProposalWrapper(
+            entityUrn=flow_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(database_key.platform),
+                instance=make_dataplatform_instance_urn(
+                    platform=database_key.platform,
+                    instance=database_key.instance,
+                ),
+            ),
+        ).as_workunit()
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=flow_urn,
+        aspect=ContainerClass(container=database_key.as_urn()),
+    ).as_workunit()
+
+
+def _get_procedure_flow_name(
+    database_key: DatabaseKey, schema_key: Optional[SchemaKey]
+) -> str:
+    if schema_key:
+        procedure_flow_name = (
+            f"{schema_key.database}.{schema_key.db_schema}.stored_procedures"
+        )
+    else:
+        procedure_flow_name = f"{database_key.database}.stored_procedures"
+    return procedure_flow_name
+
+
+def _generate_job_workunits(
+    database_key: DatabaseKey,
+    schema_key: Optional[SchemaKey],
+    procedure: BaseProcedure,
+) -> Iterable[MetadataWorkUnit]:
+    """Generate job workunits for database, schema and procedure"""
+
+    job_urn = procedure.to_urn(database_key, schema_key)
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=job_urn,
+        aspect=DataJobInfoClass(
+            name=procedure.name,
+            type=JobContainerSubTypes.STORED_PROCEDURE,
+            description=procedure.comment,
+            customProperties=procedure.extra_properties,
+        ),
+    ).as_workunit()
+
+    yield MetadataChangeProposalWrapper(
+        entityUrn=job_urn,
+        aspect=SubTypesClass(
+            typeNames=[JobContainerSubTypes.STORED_PROCEDURE],
+        ),
+    ).as_workunit()
+
+    if database_key.instance:
+        yield MetadataChangeProposalWrapper(
+            entityUrn=job_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(database_key.platform),
+                instance=make_dataplatform_instance_urn(
+                    platform=database_key.platform,
+                    instance=database_key.instance,
+                ),
+            ),
+        ).as_workunit()
+
+    container_key = schema_key or database_key  # database_key for 2-tier
+    yield MetadataChangeProposalWrapper(
+        entityUrn=job_urn,
+        aspect=ContainerClass(container=container_key.as_urn()),
+    ).as_workunit()
+
+    # TODO: Config whether to ingest procedure code
+    if procedure.procedure_definition:
+        yield MetadataChangeProposalWrapper(
+            entityUrn=job_urn,
+            aspect=DataTransformLogicClass(
+                transforms=[
+                    DataTransformClass(
+                        queryStatement=QueryStatementClass(
+                            value=procedure.procedure_definition,
+                            language=(
+                                QueryLanguageClass.SQL
+                                if procedure.language == "SQL"
+                                # The language field uses a pretty limited enum.
+                                # The "UNKNOWN" enum value is pretty new, so we don't want to
+                                # emit it until it has broader server-side support. As a
+                                # short-term solution, we map all languages to "SQL".
+                                # TODO: Once we've released server 1.1.0, we should change
+                                # this to be "UNKNOWN" for all languages except "SQL".
+                                else QueryLanguageClass.SQL
+                            ),
+                        ),
+                    )
+                ]
+            ),
+        ).as_workunit()
+
+
+def generate_procedure_lineage(
+    *,
+    schema_resolver: SchemaResolver,
+    procedure: BaseProcedure,
+    procedure_job_urn: str,
+    default_db: Optional[str] = None,
+    default_schema: Optional[str] = None,
+    is_temp_table: Callable[[str], bool] = lambda _: False,
+    raise_: bool = False,
+) -> Iterable[MetadataChangeProposalWrapper]:
+    if procedure.procedure_definition and procedure.language == "SQL":
+        datajob_input_output = parse_procedure_code(
+            schema_resolver=schema_resolver,
+            default_db=default_db,
+            default_schema=default_schema,
+            code=procedure.procedure_definition,
+            is_temp_table=is_temp_table,
+            raise_=raise_,
+        )
+
+        if datajob_input_output:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=procedure_job_urn,
+                aspect=datajob_input_output,
+            )
+
+
+def generate_procedure_container_workunits(
+    database_key: DatabaseKey,
+    schema_key: Optional[SchemaKey],
+) -> Iterable[MetadataWorkUnit]:
+    """Generate container workunits for database and schema"""
+
+    yield from _generate_flow_workunits(database_key, schema_key)
+
+
+def generate_procedure_workunits(
+    procedure: BaseProcedure,
+    database_key: DatabaseKey,
+    schema_key: Optional[SchemaKey],
+    schema_resolver: Optional[SchemaResolver],
+) -> Iterable[MetadataWorkUnit]:
+    yield from _generate_job_workunits(database_key, schema_key, procedure)
+
+    if schema_resolver:
+        job_urn = procedure.to_urn(database_key, schema_key)
+
+        yield from auto_workunit(
+            generate_procedure_lineage(
+                schema_resolver=schema_resolver,
+                procedure=procedure,
+                procedure_job_urn=job_urn,
+                default_db=database_key.database,
+                default_schema=schema_key.db_schema if schema_key else None,
+            )
+        )

datahub/ingestion/source/sql/{mssql/stored_procedure_lineage.py → stored_procedures/lineage.py}

@@ -1,8 +1,6 @@
 import logging
-from typing import Callable, Iterable, Optional
+from typing import Callable, Optional
 
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.source.sql.mssql.job_models import StoredProcedure
 from datahub.metadata.schema_classes import DataJobInputOutputClass
 from datahub.sql_parsing.datajob import to_datajob_input_output
 from datahub.sql_parsing.schema_resolver import SchemaResolver
@@ -26,6 +24,7 @@ def parse_procedure_code(
 ) -> Optional[DataJobInputOutputClass]:
     aggregator = SqlParsingAggregator(
         platform=schema_resolver.platform,
+        platform_instance=schema_resolver.platform_instance,
         env=schema_resolver.env,
         schema_resolver=schema_resolver,
         generate_lineage=True,
@@ -56,29 +55,3 @@ def parse_procedure_code(
         mcps=mcps,
         ignore_extra_mcps=True,
     )
-
-
-# Is procedure handling generic enough to be added to SqlParsingAggregator?
-def generate_procedure_lineage(
-    *,
-    schema_resolver: SchemaResolver,
-    procedure: StoredProcedure,
-    procedure_job_urn: str,
-    is_temp_table: Callable[[str], bool] = lambda _: False,
-    raise_: bool = False,
-) -> Iterable[MetadataChangeProposalWrapper]:
-    if procedure.code:
-        datajob_input_output = parse_procedure_code(
-            schema_resolver=schema_resolver,
-            default_db=procedure.db,
-            default_schema=procedure.schema,
-            code=procedure.code,
-            is_temp_table=is_temp_table,
-            raise_=raise_,
-        )
-
-        if datajob_input_output:
-            yield MetadataChangeProposalWrapper(
-                entityUrn=procedure_job_urn,
-                aspect=datajob_input_output,
-            )