acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/ge_data_profiler.py

@@ -120,7 +120,6 @@ SNOWFLAKE = "snowflake"
 BIGQUERY = "bigquery"
 REDSHIFT = "redshift"
 DATABRICKS = "databricks"
-TRINO = "trino"
 
 # Type names for Databricks, to match Title Case types in sqlalchemy
 ProfilerTypeMapping.INT_TYPE_NAMES.append("Integer")
@@ -206,6 +205,25 @@ def get_column_unique_count_dh_patch(self: SqlAlchemyDataset, column: str) -> in
             )
         )
         return convert_to_json_serializable(element_values.fetchone()[0])
+    elif (
+        self.engine.dialect.name.lower() == GXSqlDialect.AWSATHENA
+        or self.engine.dialect.name.lower() == GXSqlDialect.TRINO
+    ):
+        return convert_to_json_serializable(
+            self.engine.execute(
+                sa.select(sa.func.approx_distinct(sa.column(column))).select_from(
+                    self._table
+                )
+            ).scalar()
+        )
+    elif self.engine.dialect.name.lower() == DATABRICKS:
+        return convert_to_json_serializable(
+            self.engine.execute(
+                sa.select(sa.func.approx_count_distinct(sa.column(column))).select_from(
+                    self._table
+                )
+            ).scalar()
+        )
     return convert_to_json_serializable(
         self.engine.execute(
             sa.select([sa.func.count(sa.func.distinct(sa.column(column)))]).select_from(
@@ -289,7 +307,6 @@ def _is_single_row_query_method(query: Any) -> bool:
         "get_column_max",
         "get_column_mean",
         "get_column_stdev",
-        "get_column_nonnull_count",
         "get_column_unique_count",
     }
     CONSTANT_ROW_QUERY_METHODS = {
@@ -313,6 +330,7 @@ def _is_single_row_query_method(query: Any) -> bool:
 
     FIRST_PARTY_SINGLE_ROW_QUERY_METHODS = {
         "get_column_unique_count_dh_patch",
+        "_get_column_cardinality",
     }
 
     # We'll do this the inefficient way since the arrays are pretty small.
@@ -479,7 +497,20 @@ class _SingleDatasetProfiler(BasicDatasetProfilerBase):
         self, column_spec: _SingleColumnSpec, column: str
     ) -> None:
         try:
-            nonnull_count = self.dataset.get_column_nonnull_count(column)
+            # Don't use Great Expectations get_column_nonnull_count because it
+            # generates this SQL:
+            #
+            #   sum(CASE WHEN (mycolumn IN (NULL) OR mycolumn IS NULL) THEN 1 ELSE 0 END)
+            #
+            # which fails for complex types (such as Databricks maps) that don't
+            # support the IN operator.
+            nonnull_count = convert_to_json_serializable(
+                self.dataset.engine.execute(
+                    sa.select(sa.func.count(sa.column(column))).select_from(
+                        self.dataset._table
+                    )
+                ).scalar()
+            )
             column_spec.nonnull_count = nonnull_count
         except Exception as e:
             logger.debug(
@@ -734,11 +765,41 @@ class _SingleDatasetProfiler(BasicDatasetProfilerBase):
     def _get_dataset_column_distinct_value_frequencies(
         self, column_profile: DatasetFieldProfileClass, column: str
     ) -> None:
-        if self.config.include_field_distinct_value_frequencies:
+        if not self.config.include_field_distinct_value_frequencies:
+            return
+        try:
+            results = self.dataset.engine.execute(
+                sa.select(
+                    [
+                        sa.column(column),
+                        sa.func.count(sa.column(column)),
+                    ]
+                )
+                .select_from(self.dataset._table)
+                .where(sa.column(column).is_not(None))
+                .group_by(sa.column(column))
+            ).fetchall()
+
             column_profile.distinctValueFrequencies = [
-                ValueFrequencyClass(value=str(value), frequency=count)
-                for value, count in self.dataset.get_column_value_counts(column).items()
+                ValueFrequencyClass(value=str(value), frequency=int(count))
+                for value, count in results
             ]
+            # sort so output is deterministic. don't do it in SQL because not all column
+            # types are sortable in SQL (such as JSON data types on Athena/Trino).
+            column_profile.distinctValueFrequencies = sorted(
+                column_profile.distinctValueFrequencies, key=lambda x: x.value
+            )
+        except Exception as e:
+            logger.debug(
+                f"Caught exception while attempting to get distinct value frequencies for column {column}. {e}"
+            )
+
+            self.report.report_warning(
+                title="Profiling: Unable to Calculate Distinct Value Frequencies",
+                message="Distinct value frequencies for the column will not be accessible",
+                context=f"{self.dataset_name}.{column}",
+                exc=e,
+            )
 
     @_run_with_query_combiner
     def _get_dataset_column_histogram(
@@ -1173,26 +1234,34 @@ class DatahubGEProfiler:
             f"Will profile {len(requests)} table(s) with {max_workers} worker(s) - this may take a while"
         )
 
-        with PerfTimer() as timer, unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_unique_count",
-            get_column_unique_count_dh_patch,
-        ), unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_bigquery",
-            _get_column_quantiles_bigquery_patch,
-        ), unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_awsathena",
-            _get_column_quantiles_awsathena_patch,
-        ), unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_median",
-            _get_column_median_patch,
-        ), concurrent.futures.ThreadPoolExecutor(
-            max_workers=max_workers
-        ) as async_executor, SQLAlchemyQueryCombiner(
-            enabled=self.config.query_combiner_enabled,
-            catch_exceptions=self.config.catch_exceptions,
-            is_single_row_query_method=_is_single_row_query_method,
-            serial_execution_fallback_enabled=True,
-        ).activate() as query_combiner:
+        with (
+            PerfTimer() as timer,
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_unique_count",
+                get_column_unique_count_dh_patch,
+            ),
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_bigquery",
+                _get_column_quantiles_bigquery_patch,
+            ),
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_awsathena",
+                _get_column_quantiles_awsathena_patch,
+            ),
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_median",
+                _get_column_median_patch,
+            ),
+            concurrent.futures.ThreadPoolExecutor(
+                max_workers=max_workers
+            ) as async_executor,
+            SQLAlchemyQueryCombiner(
+                enabled=self.config.query_combiner_enabled,
+                catch_exceptions=self.config.catch_exceptions,
+                is_single_row_query_method=_is_single_row_query_method,
+                serial_execution_fallback_enabled=True,
+            ).activate() as query_combiner,
+        ):
             # Submit the profiling requests to the thread pool executor.
             async_profiles = collections.deque(
                 async_executor.submit(
@@ -1395,12 +1464,12 @@ class DatahubGEProfiler:
                     )
                 return None
             finally:
-                if batch is not None and self.base_engine.engine.name.upper() in [
-                    "TRINO",
-                    "AWSATHENA",
+                if batch is not None and self.base_engine.engine.name.lower() in [
+                    GXSqlDialect.TRINO,
+                    GXSqlDialect.AWSATHENA,
                 ]:
                     if (
-                        self.base_engine.engine.name.upper() == "TRINO"
+                        self.base_engine.engine.name.lower() == GXSqlDialect.TRINO
                         or temp_view is not None
                     ):
                         self._drop_temp_table(batch)
@@ -1449,9 +1518,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/ge_profiling_config.py

@@ -1,12 +1,12 @@
 import datetime
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Annotated, Any, Dict, List, Optional
 
 import pydantic
 from pydantic.fields import Field
 
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel, SupportedSources
 from datahub.ingestion.source_config.operation_config import OperationConfig
 
 _PROFILING_FLAGS_TO_REPORT = {
@@ -120,28 +120,37 @@ class GEProfilingConfig(GEProfilingBaseConfig):
         "number of columns to profile goes up.",
     )
 
-    profile_if_updated_since_days: Optional[pydantic.PositiveFloat] = Field(
+    profile_if_updated_since_days: Annotated[
+        Optional[pydantic.PositiveFloat], SupportedSources(["snowflake", "bigquery"])
+    ] = Field(
         default=None,
         description="Profile table only if it has been updated since these many number of days. "
         "If set to `null`, no constraint of last modified time for tables to profile. "
         "Supported only in `snowflake` and `BigQuery`.",
     )
 
-    profile_table_size_limit: Optional[int] = Field(
+    profile_table_size_limit: Annotated[
+        Optional[int],
+        SupportedSources(["snowflake", "bigquery", "unity-catalog", "oracle"]),
+    ] = Field(
         default=5,
         description="Profile tables only if their size is less than specified GBs. If set to `null`, "
         "no limit on the size of tables to profile. Supported only in `Snowflake`, `BigQuery` and "
         "`Databricks`. Supported for `Oracle` based on calculated size from gathered stats.",
     )
 
-    profile_table_row_limit: Optional[int] = Field(
+    profile_table_row_limit: Annotated[
+        Optional[int], SupportedSources(["snowflake", "bigquery", "oracle"])
+    ] = Field(
         default=5000000,
         description="Profile tables only if their row count is less than specified count. "
         "If set to `null`, no limit on the row count of tables to profile. Supported only in "
         "`Snowflake`, `BigQuery`. Supported for `Oracle` based on gathered stats.",
     )
 
-    profile_table_row_count_estimate_only: bool = Field(
+    profile_table_row_count_estimate_only: Annotated[
+        bool, SupportedSources(["postgres", "mysql"])
+    ] = Field(
         default=False,
         description="Use an approximate query for row count. This will be much faster but slightly "
         "less accurate. Only supported for Postgres and MySQL. ",
@@ -157,29 +166,35 @@ class GEProfilingConfig(GEProfilingBaseConfig):
     # Hidden option - used for debugging purposes.
     catch_exceptions: bool = Field(default=True, description="")
 
-    partition_profiling_enabled: bool = Field(
+    partition_profiling_enabled: Annotated[
+        bool, SupportedSources(["athena", "bigquery"])
+    ] = Field(
         default=True,
         description="Whether to profile partitioned tables. Only BigQuery and Aws Athena supports this. "
         "If enabled, latest partition data is used for profiling.",
     )
-    partition_datetime: Optional[datetime.datetime] = Field(
+    partition_datetime: Annotated[
+        Optional[datetime.datetime], SupportedSources(["bigquery"])
+    ] = Field(
         default=None,
         description="If specified, profile only the partition which matches this datetime. "
         "If not specified, profile the latest partition. Only Bigquery supports this.",
     )
-    use_sampling: bool = Field(
+    use_sampling: Annotated[bool, SupportedSources(["bigquery", "snowflake"])] = Field(
         default=True,
         description="Whether to profile column level stats on sample of table. Only BigQuery and Snowflake support this. "
         "If enabled, profiling is done on rows sampled from table. Sampling is not done for smaller tables. ",
     )
 
-    sample_size: int = Field(
+    sample_size: Annotated[int, SupportedSources(["bigquery", "snowflake"])] = Field(
         default=10000,
         description="Number of rows to be sampled from table for column level profiling."
         "Applicable only if `use_sampling` is set to True.",
     )
 
-    profile_external_tables: bool = Field(
+    profile_external_tables: Annotated[
+        bool, SupportedSources(["redshift", "snowflake"])
+    ] = Field(
         default=False,
         description="Whether to profile external tables. Only Snowflake and Redshift supports this.",
     )

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