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

datahub/ingestion/source/abs/source.py

@@ -44,7 +44,10 @@ from datahub.ingestion.source.azure.abs_utils import (
     get_key_prefix,
     strip_abs_prefix,
 )
-from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
+from datahub.ingestion.source.data_lake_common.data_lake_utils import (
+    ContainerWUCreator,
+    add_partition_columns_to_schema,
+)
 from datahub.ingestion.source.schema_inference import avro, csv_tsv, json, parquet
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -53,10 +56,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
-    SchemaField,
-    SchemaFieldDataType,
     SchemaMetadata,
-    StringTypeClass,
 )
 from datahub.metadata.schema_classes import (
     DataPlatformInstanceClass,
@@ -223,36 +223,12 @@ class ABSSource(StatefulIngestionSourceBase):
         fields = sorted(fields, key=lambda f: f.fieldPath)
 
         if self.source_config.add_partition_columns_to_schema:
-            self.add_partition_columns_to_schema(
+            add_partition_columns_to_schema(
                 fields=fields, path_spec=path_spec, full_path=table_data.full_path
             )
 
         return fields
 
-    def add_partition_columns_to_schema(
-        self, path_spec: PathSpec, full_path: str, fields: List[SchemaField]
-    ) -> None:
-        vars = path_spec.get_named_vars(full_path)
-        if vars is not None and "partition" in vars:
-            for partition in vars["partition"].values():
-                partition_arr = partition.split("=")
-                if len(partition_arr) != 2:
-                    logger.debug(
-                        f"Could not derive partition key from partition field {partition}"
-                    )
-                    continue
-                partition_key = partition_arr[0]
-                fields.append(
-                    SchemaField(
-                        fieldPath=f"{partition_key}",
-                        nativeDataType="string",
-                        type=SchemaFieldDataType(StringTypeClass()),
-                        isPartitioningKey=True,
-                        nullable=True,
-                        recursive=False,
-                    )
-                )
-
     def _create_table_operation_aspect(self, table_data: TableData) -> OperationClass:
         reported_time = int(time.time() * 1000)
 

datahub/ingestion/source/aws/glue.py

@@ -75,6 +75,7 @@ from datahub.ingestion.source.aws.tag_entities import (
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.glue_profiling_config import GlueProfilingConfig
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
@@ -275,6 +276,13 @@ class GlueSourceReport(StaleEntityRemovalSourceReport):
 @capability(
     SourceCapability.LINEAGE_FINE, "Support via the `emit_s3_lineage` config field"
 )
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+    ],
+)
 class GlueSource(StatefulIngestionSourceBase):
     """
     Note: if you also have files in S3 that you'd like to ingest, we recommend you use Glue's built-in data catalog. See [here](../../../../docs/generated/ingestion/sources/s3.md) for a quick guide on how to set up a crawler on Glue and ingest the outputs with DataHub.

datahub/ingestion/source/cassandra/cassandra.py

@@ -296,13 +296,11 @@ class CassandraSource(StatefulIngestionSourceBase):
             qualified_name=dataset_name,
             description=view.comment,
             custom_properties=self._get_dataset_custom_props(view),
-            extra_aspects=[
-                ViewPropertiesClass(
-                    materialized=True,
-                    viewLogic=view.where_clause,  # Use the WHERE clause as view logic
-                    viewLanguage="CQL",  # Use "CQL" as the language
-                ),
-            ],
+            view_definition=ViewPropertiesClass(
+                materialized=True,
+                viewLogic=view.where_clause,  # Use the WHERE clause as view logic
+                viewLanguage="CQL",  # Use "CQL" as the language
+            ),
         )
 
         # Construct and emit lineage off of 'base_table_name'

datahub/ingestion/source/common/subtypes.py

@@ -69,6 +69,8 @@ class BIContainerSubTypes(StrEnum):
     SIGMA_WORKSPACE = "Sigma Workspace"
     SIGMA_WORKBOOK = "Sigma Workbook"
     MODE_COLLECTION = "Collection"
+    GRAFANA_FOLDER = "Folder"
+    GRAFANA_DASHBOARD = "Dashboard"
 
 
 class FlowContainerSubTypes(StrEnum):

datahub/ingestion/source/data_lake_common/data_lake_utils.py

@@ -25,10 +25,16 @@ from datahub.ingestion.source.data_lake_common.object_store import (
     get_object_store_bucket_name,
     get_object_store_for_uri,
 )
+from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
 from datahub.ingestion.source.gcs.gcs_utils import (
     get_gcs_prefix,
     is_gcs_uri,
 )
+from datahub.metadata.schema_classes import (
+    SchemaFieldClass,
+    SchemaFieldDataTypeClass,
+    StringTypeClass,
+)
 
 # hide annoying debug errors from py4j
 logging.getLogger("py4j").setLevel(logging.ERROR)
@@ -39,6 +45,37 @@ PLATFORM_GCS = "gcs"
 PLATFORM_ABS = "abs"
 
 
+def add_partition_columns_to_schema(
+    path_spec: PathSpec, full_path: str, fields: List[SchemaFieldClass]
+) -> None:
+    # Check if using fieldPath v2 format
+    is_fieldpath_v2 = any(
+        field.fieldPath.startswith("[version=2.0]") for field in fields
+    )
+
+    # Extract partition information from path
+    partition_keys = path_spec.get_partition_from_path(full_path)
+    if not partition_keys:
+        return
+
+    # Add partition fields to schema
+    for partition_key in partition_keys:
+        fields.append(
+            SchemaFieldClass(
+                fieldPath=(
+                    f"{partition_key[0]}"
+                    if not is_fieldpath_v2
+                    else f"[version=2.0].[type=string].{partition_key[0]}"
+                ),
+                nativeDataType="string",
+                type=SchemaFieldDataTypeClass(StringTypeClass()),
+                isPartitioningKey=True,
+                nullable=False,
+                recursive=False,
+            )
+        )
+
+
 class ContainerWUCreator:
     processed_containers: List[str]
 

datahub/ingestion/source/datahub/datahub_source.py

@@ -6,7 +6,9 @@ from typing import Dict, Iterable, List, Optional
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
+    SourceCapability,
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
@@ -37,6 +39,7 @@ logger = logging.getLogger(__name__)
 @platform_name("DataHub")
 @config_class(DataHubSourceConfig)
 @support_status(SupportStatus.TESTING)
+@capability(SourceCapability.CONTAINERS, "Enabled by default")
 class DataHubSource(StatefulIngestionSourceBase):
     platform: str = "datahub"
 

datahub/ingestion/source/dbt/dbt_common.py

@@ -145,6 +145,9 @@ class DBTSourceReport(StaleEntityRemovalSourceReport):
 
     nodes_filtered: LossyList[str] = field(default_factory=LossyList)
 
+    duplicate_sources_dropped: Optional[int] = None
+    duplicate_sources_references_updated: Optional[int] = None
+
 
 class EmitDirective(ConfigEnum):
     """A holder for directives for emission for specific types of entities"""
@@ -370,6 +373,12 @@ class DBTCommonConfig(
         "Set to False to skip it for engines like AWS Athena where it's not required.",
     )
 
+    drop_duplicate_sources: bool = Field(
+        default=True,
+        description="When enabled, drops sources that have the same name in the target platform as a model. "
+        "This ensures that lineage is generated reliably, but will lose any documentation associated only with the source.",
+    )
+
     @validator("target_platform")
     def validate_target_platform_value(cls, target_platform: str) -> str:
         if target_platform.lower() == DBT_PLATFORM:
@@ -509,7 +518,7 @@ class DBTNode:
     raw_code: Optional[str]
 
     dbt_adapter: str
-    dbt_name: str
+    dbt_name: str  # dbt unique identifier
     dbt_file_path: Optional[str]
     dbt_package_name: Optional[str]  # this is pretty much always present
 
@@ -975,6 +984,8 @@ class DBTSourceBase(StatefulIngestionSourceBase):
         self._infer_schemas_and_update_cll(all_nodes_map)
 
         nodes = self._filter_nodes(all_nodes)
+        nodes = self._drop_duplicate_sources(nodes)
+
         non_test_nodes = [
             dataset_node for dataset_node in nodes if dataset_node.node_type != "test"
         ]
@@ -1000,7 +1011,7 @@ class DBTSourceBase(StatefulIngestionSourceBase):
         return self.config.node_name_pattern.allowed(key)
 
     def _filter_nodes(self, all_nodes: List[DBTNode]) -> List[DBTNode]:
-        nodes = []
+        nodes: List[DBTNode] = []
         for node in all_nodes:
             key = node.dbt_name
 
@@ -1012,6 +1023,62 @@ class DBTSourceBase(StatefulIngestionSourceBase):
 
         return nodes
 
+    def _drop_duplicate_sources(self, original_nodes: List[DBTNode]) -> List[DBTNode]:
+        """Detect and correct cases where a model and source have the same name.
+
+        In these cases, we don't want to generate both because they'll have the same
+        urn and hence overwrite each other. Instead, we drop the source and update
+        references to it to point at the model.
+
+        The risk here is that the source might have documentation that'd be lost,
+        which is why we maintain optionality with a config flag.
+        """
+        if not self.config.drop_duplicate_sources:
+            return original_nodes
+
+        self.report.duplicate_sources_dropped = 0
+        self.report.duplicate_sources_references_updated = 0
+
+        # Pass 1 - find all model names in the warehouse.
+        warehouse_model_names: Dict[str, str] = {}  # warehouse name -> model unique id
+        for node in original_nodes:
+            if node.node_type == "model" and node.exists_in_target_platform:
+                warehouse_model_names[node.get_db_fqn()] = node.dbt_name
+
+        # Pass 2 - identify + drop duplicate sources.
+        source_references_to_update: Dict[
+            str, str
+        ] = {}  # source unique id -> model unique id
+        nodes: List[DBTNode] = []
+        for node in original_nodes:
+            if (
+                node.node_type == "source"
+                and node.exists_in_target_platform
+                and (model_name := warehouse_model_names.get(node.get_db_fqn()))
+            ):
+                self.report.warning(
+                    title="Duplicate model and source names detected",
+                    message="We found a dbt model and dbt source with the same name. To ensure reliable lineage generation, the source node was ignored. "
+                    "If you associated documentation/tags/other metadata with the source, it will be lost. "
+                    "To avoid this, you should remove the source node from your dbt project and replace any `source(<source_name>)` calls with `ref(<model_name>)`.",
+                    context=f"{node.dbt_name} (called {node.get_db_fqn()} in {self.config.target_platform}) duplicates {model_name}",
+                )
+                self.report.duplicate_sources_dropped += 1
+                source_references_to_update[node.dbt_name] = model_name
+            else:
+                nodes.append(node)
+
+        # Pass 3 - update references to the dropped sources.
+        for node in nodes:
+            for i, current_upstream in enumerate(node.upstream_nodes):
+                if current_upstream in source_references_to_update:
+                    node.upstream_nodes[i] = source_references_to_update[
+                        current_upstream
+                    ]
+                    self.report.duplicate_sources_references_updated += 1
+
+        return nodes
+
     @staticmethod
     def _to_schema_info(schema_fields: List[SchemaField]) -> SchemaInfo:
         return {column.fieldPath: column.nativeDataType for column in schema_fields}

datahub/ingestion/source/delta_lake/source.py

@@ -85,6 +85,7 @@ OPERATION_STATEMENT_TYPES = {
 @config_class(DeltaLakeSourceConfig)
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.TAGS, "Can extract S3 object/bucket tags if enabled")
+@capability(SourceCapability.CONTAINERS, "Enabled by default")
 class DeltaLakeSource(StatefulIngestionSourceBase):
     """
     This plugin extracts:

datahub/ingestion/source/ge_data_profiler.py

@@ -1497,9 +1497,17 @@ class DatahubGEProfiler:
                 logger.error(
                     f"Unexpected {pretty_name} while profiling. Should have 3 parts but has {len(name_parts)} parts."
                 )
+            if platform == DATABRICKS:
+                # TODO: Review logic for BigQuery as well, probably project.dataset.table should be quoted there as well
+                quoted_name = ".".join(
+                    batch.engine.dialect.identifier_preparer.quote(part)
+                    for part in name_parts
+                )
+                batch._table = sa.text(quoted_name)
+                logger.debug(f"Setting quoted table name to be {batch._table}")
             # If we only have two parts that means the project_id is missing from the table name and we add it
             # Temp tables has 3 parts while normal tables only has 2 parts
-            if len(str(batch._table).split(".")) == 2:
+            elif len(str(batch._table).split(".")) == 2:
                 batch._table = sa.text(f"{name_parts[0]}.{str(batch._table)}")
                 logger.debug(f"Setting table name to be {batch._table}")
 

datahub/ingestion/source/grafana/entity_mcp_builder.py

@@ -0,0 +1,272 @@
+from typing import Dict, List, Optional, Tuple
+
+from datahub.emitter.mce_builder import (
+    make_chart_urn,
+    make_dashboard_urn,
+    make_data_platform_urn,
+    make_dataplatform_instance_urn,
+    make_dataset_urn_with_platform_instance,
+    make_tag_urn,
+    make_user_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.source.grafana.models import Dashboard, Panel
+from datahub.ingestion.source.grafana.types import CHART_TYPE_MAPPINGS
+from datahub.metadata.schema_classes import (
+    ChangeAuditStampsClass,
+    ChartInfoClass,
+    DashboardInfoClass,
+    DataPlatformInstanceClass,
+    GlobalTagsClass,
+    OwnerClass,
+    OwnershipClass,
+    OwnershipTypeClass,
+    StatusClass,
+    TagAssociationClass,
+)
+
+
+def build_chart_mcps(
+    panel: Panel,
+    dashboard: Dashboard,
+    platform: str,
+    platform_instance: Optional[str],
+    env: str,
+    base_url: str,
+    ingest_tags: bool,
+) -> Tuple[Optional[str], str, List[MetadataChangeProposalWrapper]]:
+    """Build chart metadata change proposals"""
+    ds_urn = None
+    mcps = []
+
+    chart_urn = make_chart_urn(
+        platform,
+        f"{dashboard.uid}.{panel.id}",
+        platform_instance,
+    )
+
+    # Platform instance aspect
+    mcps.append(
+        MetadataChangeProposalWrapper(
+            entityUrn=chart_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(platform),
+                instance=make_dataplatform_instance_urn(
+                    platform=platform,
+                    instance=platform_instance,
+                )
+                if platform_instance
+                else None,
+            ),
+        )
+    )
+
+    # Status aspect
+    mcps.append(
+        MetadataChangeProposalWrapper(
+            entityUrn=chart_urn,
+            aspect=StatusClass(removed=False),
+        )
+    )
+
+    # Get input datasets
+    input_datasets = []
+    if panel.datasource_ref:
+        ds_type = panel.datasource_ref.type or "unknown"
+        ds_uid = panel.datasource_ref.uid or "unknown"
+
+        # Add Grafana dataset
+        dataset_name = f"{ds_type}.{ds_uid}.{panel.id}"
+        ds_urn = make_dataset_urn_with_platform_instance(
+            platform=platform,
+            name=dataset_name,
+            platform_instance=platform_instance,
+            env=env,
+        )
+        input_datasets.append(ds_urn)
+
+    # Chart info aspect
+    title = panel.title or f"Panel {panel.id}"
+    mcps.append(
+        MetadataChangeProposalWrapper(
+            entityUrn=chart_urn,
+            aspect=ChartInfoClass(
+                type=CHART_TYPE_MAPPINGS.get(panel.type) if panel.type else None,
+                description=panel.description,
+                title=title,
+                lastModified=ChangeAuditStampsClass(),
+                chartUrl=f"{base_url}/d/{dashboard.uid}?viewPanel={panel.id}",
+                customProperties=_build_custom_properties(panel),
+                inputs=input_datasets,
+            ),
+        )
+    )
+
+    # Tags aspect
+    if dashboard.tags and ingest_tags:
+        tags = []
+        for tag in dashboard.tags:
+            if ":" in tag:
+                key, value = tag.split(":", 1)
+                tag_urn = make_tag_urn(f"{key}.{value}")
+            else:
+                tag_urn = make_tag_urn(tag)
+            tags.append(TagAssociationClass(tag=tag_urn))
+
+        if tags:
+            mcps.append(
+                MetadataChangeProposalWrapper(
+                    entityUrn=chart_urn,
+                    aspect=GlobalTagsClass(tags=tags),
+                )
+            )
+
+    return ds_urn, chart_urn, mcps
+
+
+def build_dashboard_mcps(
+    dashboard: Dashboard,
+    platform: str,
+    platform_instance: Optional[str],
+    chart_urns: List[str],
+    base_url: str,
+    ingest_owners: bool,
+    ingest_tags: bool,
+) -> Tuple[str, List[MetadataChangeProposalWrapper]]:
+    """Build dashboard metadata change proposals"""
+    mcps = []
+    dashboard_urn = make_dashboard_urn(platform, dashboard.uid, platform_instance)
+
+    # Platform instance aspect
+    mcps.append(
+        MetadataChangeProposalWrapper(
+            entityUrn=dashboard_urn,
+            aspect=DataPlatformInstanceClass(
+                platform=make_data_platform_urn(platform),
+                instance=make_dataplatform_instance_urn(
+                    platform=platform,
+                    instance=platform_instance,
+                )
+                if platform_instance
+                else None,
+            ),
+        )
+    )
+
+    # Dashboard info aspect
+    mcps.append(
+        MetadataChangeProposalWrapper(
+            entityUrn=dashboard_urn,
+            aspect=DashboardInfoClass(
+                description=dashboard.description,
+                title=dashboard.title,
+                charts=chart_urns,
+                lastModified=ChangeAuditStampsClass(),
+                dashboardUrl=f"{base_url}/d/{dashboard.uid}",
+                customProperties=_build_dashboard_properties(dashboard),
+            ),
+        )
+    )
+
+    # Ownership aspect
+    if dashboard.uid and ingest_owners:
+        owner = _build_ownership(dashboard)
+        if owner:
+            mcps.append(
+                MetadataChangeProposalWrapper(
+                    entityUrn=dashboard_urn,
+                    aspect=owner,
+                )
+            )
+
+    # Tags aspect
+    if dashboard.tags and ingest_tags:
+        tags = [TagAssociationClass(tag=make_tag_urn(tag)) for tag in dashboard.tags]
+        if tags:
+            mcps.append(
+                MetadataChangeProposalWrapper(
+                    entityUrn=dashboard_urn,
+                    aspect=GlobalTagsClass(tags=tags),
+                )
+            )
+
+    # Status aspect
+    mcps.append(
+        MetadataChangeProposalWrapper(
+            entityUrn=dashboard_urn,
+            aspect=StatusClass(removed=False),
+        )
+    )
+
+    return dashboard_urn, mcps
+
+
+def _build_custom_properties(panel: Panel) -> Dict[str, str]:
+    """Build custom properties for chart"""
+    props = {}
+
+    if panel.type:
+        props["type"] = panel.type
+
+    if panel.datasource_ref:
+        props["datasourceType"] = panel.datasource_ref.type or ""
+        props["datasourceUid"] = panel.datasource_ref.uid or ""
+
+    for key in [
+        "description",
+        "format",
+        "pluginVersion",
+        "repeatDirection",
+        "maxDataPoints",
+    ]:
+        value = getattr(panel, key, None)
+        if value:
+            props[key] = str(value)
+
+    if panel.query_targets:
+        props["targetsCount"] = str(len(panel.query_targets))
+
+    return props
+
+
+def _build_dashboard_properties(dashboard: Dashboard) -> Dict[str, str]:
+    """Build custom properties for dashboard"""
+    props = {}
+
+    if dashboard.timezone:
+        props["timezone"] = dashboard.timezone
+
+    if dashboard.schema_version:
+        props["schema_version"] = dashboard.schema_version
+
+    if dashboard.version:
+        props["version"] = dashboard.version
+
+    if dashboard.refresh:
+        props["refresh"] = dashboard.refresh
+
+    return props
+
+
+def _build_ownership(dashboard: Dashboard) -> Optional[OwnershipClass]:
+    """Build ownership information"""
+    owners = []
+
+    if dashboard.uid:
+        owners.append(
+            OwnerClass(
+                owner=make_user_urn(dashboard.uid),
+                type=OwnershipTypeClass.TECHNICAL_OWNER,
+            )
+        )
+
+    if dashboard.created_by:
+        owner_id = dashboard.created_by.split("@")[0]
+        owners.append(
+            OwnerClass(
+                owner=make_user_urn(owner_id),
+                type=OwnershipTypeClass.DATAOWNER,
+            )
+        )
+
+    return OwnershipClass(owners=owners) if owners else None