Merge branch 'datahub-project:master' into master

Author: asikowitz

metadata-ingestion/src/datahub/emitter/mcp.py

@@ -1,6 +1,6 @@
 import dataclasses
 import json
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
 
 from datahub.emitter.aspect import ASPECT_MAP, JSON_CONTENT_TYPE
 from datahub.emitter.serialization_helper import post_json_transform, pre_json_transform
@@ -69,6 +69,7 @@ class MetadataChangeProposalWrapper:
     aspectName: Union[None, str] = None
     aspect: Union[None, _Aspect] = None
     systemMetadata: Union[None, SystemMetadataClass] = None
+    headers: Union[None, Dict[str, str]] = None
 
     def __post_init__(self) -> None:
         if self.entityUrn and self.entityType == _ENTITY_TYPE_UNSET:
@@ -112,6 +113,7 @@ def _make_mcp_without_aspects(self) -> MetadataChangeProposalClass:
             auditHeader=self.auditHeader,
             aspectName=self.aspectName,
             systemMetadata=self.systemMetadata,
+            headers=self.headers,
         )
 
     def make_mcp(self) -> MetadataChangeProposalClass:
@@ -211,6 +213,7 @@ def try_from_mcpc(
                 aspectName=mcpc.aspectName,
                 aspect=aspect,
                 systemMetadata=mcpc.systemMetadata,
+                headers=mcpc.headers,
             )
         else:
             return None
@@ -228,6 +231,7 @@ def try_from_mcl(
             changeType=mcl.changeType,
             auditHeader=mcl.auditHeader,
             systemMetadata=mcl.systemMetadata,
+            headers=mcl.headers,
         )
         return cls.try_from_mcpc(mcpc) or mcpc
 

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

@@ -301,6 +301,7 @@ class SnowflakeV2Config(
         default=AllowDenyPattern.allow_all(),
         description=(
             "List of regex patterns for structured properties to include in ingestion."
+            " Applied to tags with form `<database>.<schema>.<tag_name>`."
             " Only used if `extract_tags` and `extract_tags_as_structured_properties` are enabled."
         ),
     )

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

@@ -23,6 +23,7 @@
 from datahub.metadata.com.linkedin.pegasus2avro.structured import (
     StructuredPropertyDefinition,
 )
+from datahub.metadata.schema_classes import ChangeTypeClass
 from datahub.metadata.urns import (
     ContainerUrn,
     DatasetUrn,
@@ -81,7 +82,7 @@ def _get_tags_on_object_without_propagation(
     def create_structured_property_templates(self) -> Iterable[MetadataWorkUnit]:
         for tag in self.data_dictionary.get_all_tags():
             if not self.config.structured_property_pattern.allowed(
-                tag.tag_identifier()
+                tag._id_prefix_as_str()
             ):
                 continue
             if self.config.extract_tags_as_structured_properties:
@@ -111,6 +112,8 @@ def gen_tag_as_structured_property_workunits(
         yield MetadataChangeProposalWrapper(
             entityUrn=urn,
             aspect=aspect,
+            changeType=ChangeTypeClass.CREATE,
+            headers={"If-None-Match": "*"},
         ).as_workunit()
 
     def _get_tags_on_object_with_propagation(

metadata-ingestion/src/datahub/testing/mcp_diff.py

@@ -2,7 +2,7 @@
 import json
 import re
 from collections import defaultdict
-from typing import Any, Dict, List, Sequence, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import deepdiff.serialization
 import yaml
@@ -34,6 +34,7 @@ class AspectForDiff:
     aspect_name: str
     aspect: Dict[str, Any] = dataclasses.field(hash=False)
     delta_info: "DeltaInfo" = dataclasses.field(hash=False, repr=False)
+    headers: Optional[Dict[str, str]] = dataclasses.field(default=None, hash=False)
 
     @classmethod
     def create_from_mcp(cls, idx: int, obj: Dict[str, Any]) -> "AspectForDiff":
@@ -44,6 +45,7 @@ def create_from_mcp(cls, idx: int, obj: Dict[str, Any]) -> "AspectForDiff":
             aspect_name=obj["aspectName"],
             aspect=aspect.get("json", aspect),
             delta_info=DeltaInfo(idx=idx, original=obj),
+            headers=obj.get("headers"),
         )
 
     def __repr__(self):
@@ -240,9 +242,12 @@ def pretty(self, verbose: bool = False) -> str:
                         s.append(serialize_aspect(ga.aspect))
                 for (i, old, new), diffs in aspect_diffs.aspects_changed.items():
                     s.append(self.report_aspect(old, i, "changed") + ":")
+
+                    print_aspects = False
                     for diff_level in diffs:
                         s.append(self.report_diff_level(diff_level, i))
-                    if verbose:
+                        print_aspects |= self.is_diff_level_on_aspect(diff_level)
+                    if verbose and print_aspects:
                         s.append(f"Old aspect:\n{serialize_aspect(old.aspect)}")
                         s.append(f"New aspect:\n{serialize_aspect(new.aspect)}")
 
@@ -271,6 +276,14 @@ def report_diff_level(diff: DiffLevel, idx: int) -> str:
             f"root[{idx}].", ""
         )
 
+    @staticmethod
+    def is_diff_level_on_aspect(diff: DiffLevel) -> bool:
+        skip_print_fields = ["changeType", "headers"]
+        try:
+            return diff.path(output_format="list")[1] not in skip_print_fields
+        except IndexError:
+            return True
+
 
 def serialize_aspect(aspect: Union[AspectForDiff, Dict[str, Any]]) -> str:
     if isinstance(aspect, AspectForDiff):  # Unpack aspect

metadata-ingestion/tests/integration/snowflake/snowflake_structured_properties_golden.json

@@ -123,7 +123,7 @@
 {
     "entityType": "structuredProperty",
     "entityUrn": "urn:li:structuredProperty:snowflake.other_db.other_schema.my_other_tag",
-    "changeType": "UPSERT",
+    "changeType": "CREATE",
     "aspectName": "propertyDefinition",
     "aspect": {
         "json": {
@@ -147,6 +147,9 @@
         "lastObserved": 1615443388097,
         "runId": "snowflake-2025_01_07-13_38_56-3fo398",
         "lastRunId": "no-run-id-provided"
+    },
+    "headers": {
+        "If-None-Match": "*"
     }
 },
 {
@@ -2978,7 +2981,7 @@
 {
     "entityType": "structuredProperty",
     "entityUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.security",
-    "changeType": "UPSERT",
+    "changeType": "CREATE",
     "aspectName": "propertyDefinition",
     "aspect": {
         "json": {
@@ -3002,6 +3005,9 @@
         "lastObserved": 1615443388097,
         "runId": "snowflake-2025_01_07-13_38_56-3fo398",
         "lastRunId": "no-run-id-provided"
+    },
+    "headers": {
+        "If-None-Match": "*"
     }
 },
 {
@@ -3318,7 +3324,7 @@
 {
     "entityType": "structuredProperty",
     "entityUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.my_tag_0",
-    "changeType": "UPSERT",
+    "changeType": "CREATE",
     "aspectName": "propertyDefinition",
     "aspect": {
         "json": {
@@ -3342,12 +3348,15 @@
         "lastObserved": 1615443388097,
         "runId": "snowflake-2025_01_07-13_38_56-3fo398",
         "lastRunId": "no-run-id-provided"
+    },
+    "headers": {
+        "If-None-Match": "*"
     }
 },
 {
     "entityType": "structuredProperty",
     "entityUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.my_tag_1",
-    "changeType": "UPSERT",
+    "changeType": "CREATE",
     "aspectName": "propertyDefinition",
     "aspect": {
         "json": {
@@ -3371,35 +3380,9 @@
         "lastObserved": 1615443388097,
         "runId": "snowflake-2025_01_07-13_38_56-3fo398",
         "lastRunId": "no-run-id-provided"
-    }
-},
-{
-    "entityType": "structuredProperty",
-    "entityUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.my_tag_2",
-    "changeType": "UPSERT",
-    "aspectName": "propertyDefinition",
-    "aspect": {
-        "json": {
-            "qualifiedName": "snowflake.test_db.test_schema.my_tag_2",
-            "displayName": "my_tag_2",
-            "valueType": "urn:li:dataType:datahub.string",
-            "cardinality": "SINGLE",
-            "entityTypes": [
-                "urn:li:entityType:datahub.container",
-                "urn:li:entityType:datahub.dataset",
-                "urn:li:entityType:datahub.schemaField"
-            ],
-            "immutable": false,
-            "lastModified": {
-                "time": 1615443388097,
-                "actor": "urn:li:corpuser:datahub"
-            }
-        }
     },
-    "systemMetadata": {
-        "lastObserved": 1615443388097,
-        "runId": "snowflake-2025_01_07-13_38_56-3fo398",
-        "lastRunId": "no-run-id-provided"
+    "headers": {
+        "If-None-Match": "*"
     }
 },
 {
@@ -3666,14 +3649,6 @@
                             "string": "my_value_1"
                         }
                     ]
-                },
-                {
-                    "propertyUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.my_tag_2",
-                    "values": [
-                        {
-                            "string": "my_value_2"
-                        }
-                    ]
                 }
             ]
         }
@@ -5007,22 +4982,6 @@
         "lastRunId": "no-run-id-provided"
     }
 },
-{
-    "entityType": "structuredProperty",
-    "entityUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.my_tag_2",
-    "changeType": "UPSERT",
-    "aspectName": "status",
-    "aspect": {
-        "json": {
-            "removed": false
-        }
-    },
-    "systemMetadata": {
-        "lastObserved": 1615443388097,
-        "runId": "snowflake-2025_01_07-13_38_56-3fo398",
-        "lastRunId": "no-run-id-provided"
-    }
-},
 {
     "entityType": "structuredProperty",
     "entityUrn": "urn:li:structuredProperty:snowflake.test_db.test_schema.security",
@@ -5055,4 +5014,4 @@
         "lastRunId": "no-run-id-provided"
     }
 }
-]
+]

metadata-ingestion/tests/integration/snowflake/test_snowflake.py

@@ -209,6 +209,9 @@ def test_snowflake_tags_as_structured_properties(
                     type="snowflake",
                     config=SnowflakeV2Config(
                         extract_tags_as_structured_properties=True,
+                        structured_property_pattern=AllowDenyPattern(
+                            deny=["test_db.test_schema.my_tag_2"]
+                        ),
                         extract_tags=TagOption.without_lineage,
                         account_id="ABC12345.ap-south-1.aws",
                         username="TST_USR",

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java

@@ -121,6 +121,7 @@ public void appendRunId(
 
     // Create an upsert document that will be used if the document doesn't exist
     Map<String, Object> upsert = new HashMap<>();
+    upsert.put("urn", urn.toString());
     upsert.put("runId", Collections.singletonList(runId));
 
     esWriteDAO.applyScriptUpdate(

metadata-io/src/test/java/com/linkedin/metadata/search/ElasticSearchServiceTest.java

@@ -89,6 +89,7 @@ public void testAppendRunId_ValidRunId() {
     // Verify upsert document
     Map<String, Object> capturedUpsert = upsertCaptor.getValue();
     assertEquals(capturedUpsert.get("runId"), Collections.singletonList(runId));
+    assertEquals(capturedUpsert.get("urn"), TEST_URN.toString());
   }
 
   @Test
@@ -121,6 +122,7 @@ public void testAppendRunId_NullRunId() {
 
     Map<String, Object> capturedUpsert = upsertCaptor.getValue();
     assertEquals(capturedUpsert.get("runId"), Collections.singletonList(null));
+    assertEquals(capturedUpsert.get("urn"), TEST_URN.toString());
   }
 
   @Test(expectedExceptions = NullPointerException.class)