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

datahub/ingestion/source/sql/clickhouse.py

@@ -18,6 +18,7 @@ from sqlalchemy.sql import sqltypes
 from sqlalchemy.types import BOOLEAN, DATE, DATETIME, INTEGER
 
 import datahub.emitter.mce_builder as builder
+from datahub.configuration.common import HiddenFromDocs, LaxStr
 from datahub.configuration.source_common import DatasetLineageProviderConfigBase
 from datahub.configuration.time_window_config import BaseTimeWindowConfig
 from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated
@@ -32,6 +33,7 @@ from datahub.ingestion.api.decorators import (
     support_status,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.sql_common import (
     SqlWorkUnit,
     logger,
@@ -127,16 +129,20 @@ class ClickHouseConfig(
 ):
     # defaults
     host_port: str = Field(default="localhost:8123", description="ClickHouse host URL.")
-    scheme: str = Field(default="clickhouse", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="clickhouse")
     password: pydantic.SecretStr = Field(
         default=pydantic.SecretStr(""), description="password"
     )
-    secure: Optional[bool] = Field(default=None, description="")
-    protocol: Optional[str] = Field(default=None, description="")
+    secure: Optional[bool] = Field(
+        default=None, description="[deprecated] Use uri_opts instead."
+    )
+    protocol: Optional[str] = Field(
+        default=None, description="[deprecated] Use uri_opts instead."
+    )
     _deprecate_secure = pydantic_field_deprecated("secure")
     _deprecate_protocol = pydantic_field_deprecated("protocol")
 
-    uri_opts: Dict[str, str] = Field(
+    uri_opts: Dict[str, LaxStr] = Field(
         default={},
         description="The part of the URI and it's used to provide additional configuration options or parameters for the database connection.",
     )
@@ -145,7 +151,11 @@ class ClickHouseConfig(
     )
     include_materialized_views: Optional[bool] = Field(default=True, description="")
 
-    def get_sql_alchemy_url(self, current_db=None):
+    def get_sql_alchemy_url(
+        self,
+        uri_opts: Optional[Dict[str, Any]] = None,
+        current_db: Optional[str] = None,
+    ) -> str:
         url = make_url(
             super().get_sql_alchemy_url(uri_opts=self.uri_opts, current_db=current_db)
         )
@@ -180,9 +190,9 @@ class ClickHouseConfig(
                 "Initializing uri_opts from deprecated secure or protocol options"
             )
             values["uri_opts"] = {}
-            if secure:
-                values["uri_opts"]["secure"] = secure
-            if protocol:
+            if secure is not None:
+                values["uri_opts"]["secure"] = str(secure)
+            if protocol is not None:
                 values["uri_opts"]["protocol"] = protocol
             logger.debug(f"uri_opts: {uri_opts}")
         elif (secure or protocol) and uri_opts:
@@ -375,8 +385,18 @@ clickhouse_datetime_format = "%Y-%m-%d %H:%M:%S"
 @platform_name("ClickHouse")
 @config_class(ClickHouseConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Enabled by default to get lineage for views via `include_view_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.TABLE,
+    ],
+)
 class ClickHouseSource(TwoTierSQLAlchemySource):
     """
     This plugin extracts the following:

datahub/ingestion/source/sql/cockroachdb.py

@@ -1,6 +1,6 @@
 from pydantic.fields import Field
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -14,8 +14,10 @@ from datahub.ingestion.source.sql.postgres import PostgresConfig, PostgresSource
 
 
 class CockroachDBConfig(PostgresConfig):
-    scheme = Field(default="cockroachdb+psycopg2", description="database scheme")
-    schema_pattern = Field(
+    scheme: HiddenFromDocs[str] = Field(
+        default="cockroachdb+psycopg2", description="database scheme"
+    )
+    schema_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern(deny=["information_schema", "crdb_internal"])
     )
 
@@ -26,7 +28,6 @@ class CockroachDBConfig(PostgresConfig):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class CockroachDBSource(PostgresSource):
     config: CockroachDBConfig
 

datahub/ingestion/source/sql/druid.py

@@ -1,10 +1,12 @@
 # This import verifies that the dependencies are available.
+from typing import Any, Dict, Optional
+
 import pydruid  # noqa: F401
 from pydantic.fields import Field
 from pydruid.db.sqlalchemy import DruidDialect
 from sqlalchemy.exc import ResourceClosedError
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.ingestion.api.decorators import (
     SourceCapability,
     SupportStatus,
@@ -32,14 +34,17 @@ DruidDialect.get_table_names = get_table_names
 
 class DruidConfig(BasicSQLAlchemyConfig):
     # defaults
-    scheme: str = "druid"
+    scheme: HiddenFromDocs[str] = "druid"
     schema_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern(deny=["^(lookup|sysgit|view).*"]),
         description="regex patterns for schemas to filter in ingestion.",
     )
 
-    def get_sql_alchemy_url(self):
-        return f"{super().get_sql_alchemy_url()}/druid/v2/sql/"
+    def get_sql_alchemy_url(
+        self, uri_opts: Optional[Dict[str, Any]] = None, database: Optional[str] = None
+    ) -> str:
+        base_url = super().get_sql_alchemy_url(uri_opts=uri_opts, database=database)
+        return f"{base_url}/druid/v2/sql/"
 
     """
     The pydruid library already formats the table name correctly, so we do not

datahub/ingestion/source/sql/hana.py

@@ -27,7 +27,9 @@ class HanaConfig(BasicSQLAlchemyConfig):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 class HanaSource(SQLAlchemySource):
     def __init__(self, config: HanaConfig, ctx: PipelineContext):
         super().__init__(config, ctx, "hana")

datahub/ingestion/source/sql/hive.py

@@ -6,7 +6,7 @@ from enum import Enum
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
-from pydantic.class_validators import validator
+from pydantic import validator
 from pydantic.fields import Field
 
 # This import verifies that the dependencies are available.
@@ -14,6 +14,7 @@ from pyhive import hive  # noqa: F401
 from pyhive.sqlalchemy_hive import HiveDate, HiveDecimal, HiveDialect, HiveTimestamp
 from sqlalchemy.engine.reflection import Inspector
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
@@ -139,7 +140,7 @@ class StoragePathParser:
                 path = f"{parsed.netloc}/{parsed.path.lstrip('/')}"
 
             elif platform == StoragePlatform.AZURE:
-                if scheme in ("abfs", "abfss"):
+                if scheme in ("abfs", "abfss", "wasbs"):
                     # Format: abfss://container@account.dfs.core.windows.net/path
                     container = parsed.netloc.split("@")[0]
                     path = f"{container}/{parsed.path.lstrip('/')}"
@@ -153,7 +154,7 @@ class StoragePathParser:
 
             elif platform == StoragePlatform.DBFS:
                 # For DBFS, use path as-is
-                path = parsed.path.lstrip("/")
+                path = "/" + parsed.path.lstrip("/")
 
             elif platform == StoragePlatform.LOCAL:
                 # For local files, use full path
@@ -169,7 +170,6 @@ class StoragePathParser:
             # Clean up the path
             path = path.rstrip("/")  # Remove trailing slashes
             path = re.sub(r"/+", "/", path)  # Normalize multiple slashes
-            path = f"/{path}"
 
             return platform, path
 
@@ -637,8 +637,13 @@ def get_view_definition_patched(self, connection, view_name, schema=None, **kw):
             self.identifier_preparer.quote_identifier(schema),
             self.identifier_preparer.quote_identifier(view_name),
         )
-    row = connection.execute(f"SHOW CREATE TABLE {full_table}").fetchone()
-    return row[0]
+    # Hive responds to the SHOW CREATE TABLE with the full view DDL,
+    # including the view definition. However, for multiline view definitions,
+    # it returns multiple rows (of one column each), each with a part of the definition.
+    # Any whitespace at the beginning/end of each view definition line is lost.
+    rows = connection.execute(f"SHOW CREATE TABLE {full_table}").fetchall()
+    parts = [row[0] for row in rows]
+    return "\n".join(parts)
 
 
 HiveDialect.get_view_names = get_view_names_patched
@@ -647,10 +652,10 @@ HiveDialect.get_view_definition = get_view_definition_patched
 
 class HiveConfig(TwoTierSQLAlchemyConfig):
     # defaults
-    scheme: str = Field(default="hive", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="hive")
 
     # Overriding as table location lineage is richer implementation here than with include_table_location_lineage
-    include_table_location_lineage: bool = Field(default=False, hidden_from_docs=True)
+    include_table_location_lineage: HiddenFromDocs[bool] = Field(default=False)
 
     emit_storage_lineage: bool = Field(
         default=False,
@@ -862,3 +867,18 @@ class HiveSource(TwoTierSQLAlchemySource):
                 return partition_column.get("column_names")
 
         return []
+
+    def get_table_properties(
+        self, inspector: Inspector, schema: str, table: str
+    ) -> Tuple[Optional[str], Dict[str, str], Optional[str]]:
+        (description, properties, location) = super().get_table_properties(
+            inspector, schema, table
+        )
+
+        new_properties = {}
+        for key, value in properties.items():
+            if key and key[-1] == ":":
+                new_properties[key[:-1]] = value
+            else:
+                new_properties[key] = value
+        return (description, new_properties, location)

datahub/ingestion/source/sql/hive_metastore.py

@@ -1,17 +1,15 @@
 import base64
+import dataclasses
 import json
 import logging
 from collections import namedtuple
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
-from pydantic.dataclasses import dataclass
-from pydantic.fields import Field
-
-# This import verifies that the dependencies are available.
+from pydantic import Field
 from sqlalchemy import create_engine, text
 from sqlalchemy.engine.reflection import Inspector
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -27,6 +25,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
@@ -36,7 +35,6 @@ from datahub.ingestion.source.sql.sql_common import (
 from datahub.ingestion.source.sql.sql_config import (
     BasicSQLAlchemyConfig,
     SQLCommonConfig,
-    make_sqlalchemy_uri,
 )
 from datahub.ingestion.source.sql.sql_utils import (
     add_table_to_schema_container,
@@ -46,13 +44,13 @@ from datahub.ingestion.source.sql.sql_utils import (
     gen_schema_key,
     get_domain_wu,
 )
+from datahub.ingestion.source.sql.sqlalchemy_uri import make_sqlalchemy_uri
 from datahub.ingestion.source.state.stateful_ingestion_base import JobId
 from datahub.metadata.com.linkedin.pegasus2avro.common import StatusClass
 from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaField
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     DatasetPropertiesClass,
     SubTypesClass,
     ViewPropertiesClass,
@@ -67,13 +65,13 @@ TableKey = namedtuple("TableKey", ["schema", "table"])
 
 
 class HiveMetastoreConfigMode(StrEnum):
-    hive: str = "hive"
-    presto: str = "presto"
-    presto_on_hive: str = "presto-on-hive"
-    trino: str = "trino"
+    hive = "hive"
+    presto = "presto"
+    presto_on_hive = "presto-on-hive"
+    trino = "trino"
 
 
-@dataclass
+@dataclasses.dataclass
 class ViewDataset:
     dataset_name: str
     schema_name: str
@@ -99,7 +97,7 @@ class HiveMetastore(BasicSQLAlchemyConfig):
         default="localhost:3306",
         description="Host URL and port to connect to. Example: localhost:3306",
     )
-    scheme: str = Field(default="mysql+pymysql", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="mysql+pymysql")
 
     database_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
@@ -123,8 +121,8 @@ class HiveMetastore(BasicSQLAlchemyConfig):
         description="Dataset Subtype name to be 'Table' or 'View' Valid options: ['True', 'False']",
     )
 
-    include_view_lineage: bool = Field(
-        default=False, description="", hidden_from_docs=True
+    include_view_lineage: HiddenFromDocs[bool] = Field(
+        default=False,
     )
 
     include_catalog_name_in_ids: bool = Field(
@@ -161,12 +159,22 @@ class HiveMetastore(BasicSQLAlchemyConfig):
 @platform_name("Hive Metastore")
 @config_class(HiveMetastore)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 @capability(SourceCapability.DATA_PROFILING, "Not Supported", False)
 @capability(SourceCapability.CLASSIFICATION, "Not Supported", False)
 @capability(
     SourceCapability.LINEAGE_COARSE, "View lineage is not supported", supported=False
 )
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.CATALOG,
+        SourceCapabilityModifier.SCHEMA,
+    ],
+)
 class HiveMetastoreSource(SQLAlchemySource):
     """
     This plugin extracts the following:
@@ -599,10 +607,7 @@ class HiveMetastoreSource(SQLAlchemySource):
                 yield dpi_aspect
 
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="subTypes",
                 aspect=SubTypesClass(typeNames=[self.table_subtype]),
             ).as_workunit()
 
@@ -808,10 +813,7 @@ class HiveMetastoreSource(SQLAlchemySource):
 
             # Add views subtype
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="subTypes",
                 aspect=SubTypesClass(typeNames=[self.view_subtype]),
             ).as_workunit()
 
@@ -822,10 +824,7 @@ class HiveMetastoreSource(SQLAlchemySource):
                 viewLogic=dataset.view_definition if dataset.view_definition else "",
             )
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="viewProperties",
                 aspect=view_properties_aspect,
             ).as_workunit()
 

datahub/ingestion/source/sql/mariadb.py

@@ -15,7 +15,6 @@ from datahub.ingestion.source.sql.mysql import MySQLConfig, MySQLSource
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class MariaDBSource(MySQLSource):
     def get_platform(self):
         return "mariadb"

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,
@@ -133,7 +134,22 @@ class StoredProcedure:
 
     @property
     def escape_full_name(self) -> str:
-        return f"[{self.db}].[{self.schema}].[{self.formatted_name}]"
+        return f"[{self.db}].[{self.schema}].[{self.formatted_name}]".replace(
+            "'", r"''"
+        )
+
+    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
@@ -222,7 +238,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],