feat(snowflake): add support for external DMF assertion ingestion (#16058)

Authored-by: Rajat Singh <254366865+rajatoss@users.noreply.github.com>

Author: rajatoss

docs/assertions/snowflake/snowflake_dmfs.md

@@ -180,7 +180,7 @@ snowsql -f dmf_definitions.sql
 snowsql -f dmf_associations.sql
 ```
 
-:::NOTE
+::: NOTE
 Scheduling Data Metric Function on table incurs Serverless Credit Usage in Snowflake. Refer [Billing and Pricing](https://docs.snowflake.com/en/user-guide/data-quality-intro#billing-and-pricing) for more details.
 Please ensure you DROP Data Metric Function created via dmf_associations.sql if the assertion is no longer in use.
 :::
@@ -208,6 +208,127 @@ either via CLI or the UI visible as normal assertions.
 
 `datahub ingest -c snowflake.yml`
 
+## Ingesting External (User-Created) DMFs
+
+In addition to DataHub-created DMFs, you can also ingest results from your own custom Snowflake Data Metric Functions. "External" here means DMFs that were created directly in Snowflake without using DataHub's assertion compiler - they exist outside of DataHub's management.
+
+### Why Use External DMFs?
+
+You might want to ingest external DMFs if:
+
+- **Pre-existing DMFs**: You already have DMFs in Snowflake that were created before adopting DataHub, and you want to see their results in DataHub without recreating them
+- **Custom logic**: You need DMF logic that isn't supported by DataHub's assertion compiler (e.g., complex multi-table checks)
+- **Team workflows**: Different teams manage DMFs directly in Snowflake, but you want centralized visibility in DataHub
+- **Gradual adoption**: You want to start monitoring existing data quality checks in DataHub before fully migrating to DataHub-managed assertions
+
+### Enabling External DMF Ingestion
+
+To ingest external DMFs, add the `include_externally_managed_dmfs` flag to your Snowflake recipe:
+
+```yaml
+source:
+  type: snowflake
+  config:
+    # ... connection config ...
+
+    # Enable assertion results ingestion (required)
+    include_assertion_results: true
+
+    # Enable external DMF ingestion (new)
+    include_externally_managed_dmfs: true
+
+    # Time window for assertion results
+    start_time: "-7 days"
+```
+
+Both flags must be enabled for external DMF ingestion to work.
+
+### Requirements for External DMFs
+
+**External DMFs must return `1` for SUCCESS and `0` for FAILURE.**
+
+DataHub interprets the `VALUE` column from Snowflake's `DATA_QUALITY_MONITORING_RESULTS` table as:
+
+- `VALUE = 1` → Assertion **PASSED**
+- `VALUE = 0` → Assertion **FAILED**
+
+This is because DataHub cannot interpret arbitrary return values (e.g., "100 null rows" - is that good or bad?). You must build the pass/fail logic into your DMF.
+
+::: warning What if my DMF returns other values?
+If your DMF returns values other than 0 or 1, DataHub will mark the assertion result as **ERROR**:
+
+- `VALUE = 1` → **PASSED**
+- `VALUE = 0` → **FAILED**
+- `VALUE != 0 and VALUE != 1` (e.g., 5, 100, -1) → **ERROR**
+
+The ERROR state indicates that the DMF is not configured correctly for DataHub ingestion. You can identify these cases by:
+
+1. Checking the ingestion logs for warnings like: `DMF 'my_dmf' returned invalid value 100. Expected 1 (pass) or 0 (fail). Marking as ERROR.`
+2. Looking for assertions with ERROR status in the DataHub UI
+   :::
+
+#### Example: Writing External DMFs Correctly
+
+**WRONG** - Returns raw count (DataHub can't interpret this):
+
+```sql
+CREATE DATA METRIC FUNCTION my_null_check(ARGT TABLE(col VARCHAR))
+RETURNS NUMBER AS
+$$
+  SELECT COUNT(*) FROM ARGT WHERE col IS NULL
+$$;
+-- Returns: 0, 5, 100, etc. - DataHub can't determine pass/fail!
+```
+
+**CORRECT** - Returns 1 (pass) or 0 (fail):
+
+```sql
+CREATE DATA METRIC FUNCTION my_null_check(ARGT TABLE(col VARCHAR))
+RETURNS NUMBER AS
+$$
+  SELECT CASE WHEN COUNT(*) = 0 THEN 1 ELSE 0 END
+  FROM ARGT WHERE col IS NULL
+$$;
+-- Returns: 1 if no nulls (pass), 0 if has nulls (fail)
+```
+
+**CORRECT** - With threshold:
+
+```sql
+CREATE DATA METRIC FUNCTION my_null_check_threshold(ARGT TABLE(col VARCHAR))
+RETURNS NUMBER AS
+$$
+  SELECT CASE WHEN COUNT(*) <= 10 THEN 1 ELSE 0 END
+  FROM ARGT WHERE col IS NULL
+$$;
+-- Returns: 1 if ≤10 nulls (pass), 0 if >10 nulls (fail)
+```
+
+### How External DMFs Differ from DataHub-Created DMFs
+
+| Aspect             | DataHub-Created DMFs                               | External DMFs                             |
+| ------------------ | -------------------------------------------------- | ----------------------------------------- |
+| **Naming**         | Prefixed with `datahub__`                          | Any name                                  |
+| **Definition**     | Created via `datahub assertions compile`           | Created manually in Snowflake             |
+| **Assertion Type** | Based on YAML definition (Freshness, Volume, etc.) | CUSTOM                                    |
+| **Source**         | NATIVE (defined in DataHub)                        | EXTERNAL                                  |
+| **URN Generation** | Extracted from DMF name (`datahub__<guid>`)        | Generated from Snowflake's `REFERENCE_ID` |
+
+### How External DMFs Appear in DataHub UI
+
+External DMFs appear in DataHub with:
+
+- **Assertion Type**: CUSTOM
+- **Source**: EXTERNAL
+- **Platform Instance**: Snowflake platform instance (if configured)
+- **Description**: "External Snowflake DMF: {dmf_name}"
+- **Custom Properties**:
+  - `snowflake_dmf_name`: The DMF function name
+  - `snowflake_reference_id`: Snowflake's unique identifier for the DMF-table binding
+  - `snowflake_dmf_columns`: Comma-separated list of columns the DMF operates on
+
+You can view external DMF assertions in the **Quality** tab of the associated dataset in the DataHub UI. They will show pass/fail history alongside any DataHub-created assertions.
+
 ## Caveats
 
 - Currently, Snowflake supports at most 1000 DMF-table associations at the moment so you can not define more than 1000 assertions for snowflake.

metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_assertion.py

@@ -1,15 +1,18 @@
+import json
 import logging
 from datetime import datetime
-from typing import Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 
 from datahub.emitter.mce_builder import (
     make_assertion_urn,
     make_data_platform_urn,
     make_dataplatform_instance_urn,
+    make_schema_field_urn,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import DatahubKey
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.snowflake.snowflake_config import SnowflakeV2Config
 from datahub.ingestion.source.snowflake.snowflake_connection import SnowflakeConnection
@@ -25,18 +28,55 @@
     AssertionRunStatus,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import DataPlatformInstance
+from datahub.metadata.schema_classes import (
+    AssertionInfoClass,
+    AssertionSourceClass,
+    AssertionSourceTypeClass,
+    AssertionTypeClass,
+    CustomAssertionInfoClass,
+)
 from datahub.utilities.time import datetime_to_ts_millis
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+class SnowflakeExternalDmfKey(DatahubKey):
+    """Key for generating deterministic GUIDs for external Snowflake DMFs.
+
+    Uses Snowflake's REFERENCE_ID which uniquely identifies the
+    DMF-table-column association.
+    """
+
+    platform: str = "snowflake"
+    reference_id: str
+    instance: Optional[str] = None
+
+
 class DataQualityMonitoringResult(BaseModel):
     MEASUREMENT_TIME: datetime
     METRIC_NAME: str
     TABLE_NAME: str
     TABLE_SCHEMA: str
     TABLE_DATABASE: str
     VALUE: int
+    REFERENCE_ID: str
+    ARGUMENT_NAMES: List[str]
+
+    @field_validator("ARGUMENT_NAMES", mode="before")
+    @classmethod
+    def parse_argument_names(cls, v: object) -> List[str]:
+        """Parse ARGUMENT_NAMES from JSON string.
+
+        Snowflake returns this column as a JSON-encoded string like '["col1", "col2"]'.
+        """
+        if isinstance(v, str):
+            try:
+                parsed = json.loads(v)
+                if isinstance(parsed, list):
+                    return parsed
+            except json.JSONDecodeError:
+                logger.debug(f"Failed to parse ARGUMENT_NAMES as JSON: {v}")
+        return []
 
 
 class SnowflakeAssertionsHandler:
@@ -56,20 +96,19 @@ def __init__(
     def get_assertion_workunits(
         self, discovered_datasets: List[str]
     ) -> Iterable[MetadataWorkUnit]:
+        include_external = self.config.include_externally_managed_dmfs
+
         cur = self.connection.query(
             SnowflakeQuery.dmf_assertion_results(
                 datetime_to_ts_millis(self.config.start_time),
                 datetime_to_ts_millis(self.config.end_time),
+                include_external=include_external,
             )
         )
         for db_row in cur:
-            mcp = self._process_result_row(db_row, discovered_datasets)
-            if mcp:
-                yield mcp.as_workunit(is_primary_source=False)
-
-                if mcp.entityUrn and mcp.entityUrn not in self._urns_processed:
-                    self._urns_processed.append(mcp.entityUrn)
-                    yield self._gen_platform_instance_wu(mcp.entityUrn)
+            workunits = self._process_result_row(db_row, discovered_datasets)
+            for wu in workunits:
+                yield wu
 
     def _gen_platform_instance_wu(self, urn: str) -> MetadataWorkUnit:
         # Construct a MetadataChangeProposalWrapper object for assertion platform
@@ -87,34 +126,122 @@ def _gen_platform_instance_wu(self, urn: str) -> MetadataWorkUnit:
             ),
         ).as_workunit(is_primary_source=False)
 
+    def _generate_external_dmf_guid(self, result: DataQualityMonitoringResult) -> str:
+        """Generate a stable, deterministic GUID for external DMFs."""
+        key = SnowflakeExternalDmfKey(
+            reference_id=result.REFERENCE_ID,
+            instance=self.config.platform_instance,
+        )
+        return key.guid()
+
+    def _create_assertion_info_workunit(
+        self,
+        assertion_urn: str,
+        dataset_urn: str,
+        dmf_name: str,
+        argument_names: List[str],
+        reference_id: str,
+    ) -> MetadataWorkUnit:
+        """Create AssertionInfo for external DMFs."""
+        # Field URN is only set for single-column DMFs. Multi-column DMFs are
+        # treated as table-level assertions with columns stored in custom properties.
+        field_urn: Optional[str] = None
+        if argument_names and len(argument_names) == 1:
+            field_urn = make_schema_field_urn(dataset_urn, argument_names[0])
+
+        custom_properties: Dict[str, str] = {
+            "snowflake_dmf_name": dmf_name,
+            "snowflake_reference_id": reference_id,
+        }
+        # Store all columns in custom properties regardless of count
+        if argument_names:
+            custom_properties["snowflake_dmf_columns"] = ",".join(argument_names)
+
+        assertion_info = AssertionInfoClass(
+            type=AssertionTypeClass.CUSTOM,
+            customAssertion=CustomAssertionInfoClass(
+                type="Snowflake Data Metric Function",
+                entity=dataset_urn,
+                field=field_urn,
+            ),
+            source=AssertionSourceClass(
+                type=AssertionSourceTypeClass.EXTERNAL,
+            ),
+            description=f"External Snowflake DMF: {dmf_name}",
+            customProperties=custom_properties,
+        )
+
+        return MetadataChangeProposalWrapper(
+            entityUrn=assertion_urn,
+            aspect=assertion_info,
+        ).as_workunit(is_primary_source=False)
+
     def _process_result_row(
         self, result_row: dict, discovered_datasets: List[str]
-    ) -> Optional[MetadataChangeProposalWrapper]:
+    ) -> List[MetadataWorkUnit]:
+        """Process a single DMF result row. Returns list of workunits."""
+        workunits: List[MetadataWorkUnit] = []
+
         try:
             result = DataQualityMonitoringResult.model_validate(result_row)
-            assertion_guid = result.METRIC_NAME.split("__")[-1].lower()
-            status = bool(result.VALUE)  # 1 if PASS, 0 if FAIL
+
+            is_datahub_dmf = result.METRIC_NAME.lower().startswith("datahub__")
+
+            if is_datahub_dmf:
+                assertion_guid = result.METRIC_NAME.split("__")[-1].lower()
+            else:
+                assertion_guid = self._generate_external_dmf_guid(result)
+
+            assertion_urn = make_assertion_urn(assertion_guid)
+
             assertee = self.identifiers.get_dataset_identifier(
                 result.TABLE_NAME, result.TABLE_SCHEMA, result.TABLE_DATABASE
             )
-            if assertee in discovered_datasets:
-                return MetadataChangeProposalWrapper(
-                    entityUrn=make_assertion_urn(assertion_guid),
-                    aspect=AssertionRunEvent(
-                        timestampMillis=datetime_to_ts_millis(result.MEASUREMENT_TIME),
-                        runId=result.MEASUREMENT_TIME.strftime("%Y-%m-%dT%H:%M:%SZ"),
-                        asserteeUrn=self.identifiers.gen_dataset_urn(assertee),
-                        status=AssertionRunStatus.COMPLETE,
-                        assertionUrn=make_assertion_urn(assertion_guid),
-                        result=AssertionResult(
-                            type=(
-                                AssertionResultType.SUCCESS
-                                if status
-                                else AssertionResultType.FAILURE
-                            )
-                        ),
-                    ),
+            if assertee not in discovered_datasets:
+                return []
+
+            dataset_urn = self.identifiers.gen_dataset_urn(assertee)
+
+            if result.VALUE == 1:
+                result_type = AssertionResultType.SUCCESS
+            elif result.VALUE == 0:
+                result_type = AssertionResultType.FAILURE
+            else:
+                result_type = AssertionResultType.ERROR
+                logger.warning(
+                    f"DMF '{result.METRIC_NAME}' returned invalid value {result.VALUE}. "
+                    "Expected 1 (pass) or 0 (fail). Marking as ERROR."
+                )
+
+            if not is_datahub_dmf and assertion_urn not in self._urns_processed:
+                assertion_info_wu = self._create_assertion_info_workunit(
+                    assertion_urn=assertion_urn,
+                    dataset_urn=dataset_urn,
+                    dmf_name=result.METRIC_NAME,
+                    argument_names=result.ARGUMENT_NAMES,
+                    reference_id=result.REFERENCE_ID,
                 )
+                workunits.append(assertion_info_wu)
+
+            run_event_mcp = MetadataChangeProposalWrapper(
+                entityUrn=assertion_urn,
+                aspect=AssertionRunEvent(
+                    timestampMillis=datetime_to_ts_millis(result.MEASUREMENT_TIME),
+                    runId=result.MEASUREMENT_TIME.strftime("%Y-%m-%dT%H:%M:%SZ"),
+                    asserteeUrn=dataset_urn,
+                    status=AssertionRunStatus.COMPLETE,
+                    assertionUrn=assertion_urn,
+                    result=AssertionResult(type=result_type),
+                ),
+            )
+            workunits.append(run_event_mcp.as_workunit(is_primary_source=False))
+
+            if assertion_urn not in self._urns_processed:
+                self._urns_processed.append(assertion_urn)
+                workunits.append(self._gen_platform_instance_wu(assertion_urn))
+
+            return workunits
+
         except Exception as e:
             self.report.report_warning("assertion-result-parse-failure", str(e))
-        return None
+            return []