feat(SDK) Add StructuredPropertyPatchBuilder in python sdk and provid…

…e sample CRUD files (#10824)

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/template/ArrayMergingTemplate.java

@@ -12,6 +12,8 @@
 
 public interface ArrayMergingTemplate<T extends RecordTemplate> extends Template<T> {
 
+  static final String UNIT_SEPARATOR_DELIMITER = "␟";
+
   /**
    * Takes an Array field on the {@link RecordTemplate} subtype along with a set of key fields to
    * transform into a map Avoids producing side effects by copying nodes, use resulting node and not
@@ -39,7 +41,15 @@ default JsonNode arrayFieldToMap(
                 JsonNode nodeClone = node.deepCopy();
                 if (!keyFields.isEmpty()) {
                   for (String keyField : keyFields) {
-                    String key = node.get(keyField).asText();
+                    String key;
+                    // if the keyField has a unit separator, we are working with a nested key
+                    if (keyField.contains(UNIT_SEPARATOR_DELIMITER)) {
+                      String[] keyParts = keyField.split(UNIT_SEPARATOR_DELIMITER);
+                      JsonNode keyObject = node.get(keyParts[0]);
+                      key = keyObject.get(keyParts[1]).asText();
+                    } else {
+                      key = node.get(keyField).asText();
+                    }
                     keyValue =
                         keyValue.get(key) == null
                             ? (ObjectNode) keyValue.set(key, instance.objectNode()).get(key)

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/template/AspectTemplateEngine.java

@@ -13,6 +13,7 @@
 import static com.linkedin.metadata.Constants.GLOSSARY_TERMS_ASPECT_NAME;
 import static com.linkedin.metadata.Constants.OWNERSHIP_ASPECT_NAME;
 import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTIES_ASPECT_NAME;
+import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME;
 import static com.linkedin.metadata.Constants.UPSTREAM_LINEAGE_ASPECT_NAME;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -48,6 +49,7 @@ public class AspectTemplateEngine {
               CHART_INFO_ASPECT_NAME,
               DASHBOARD_INFO_ASPECT_NAME,
               STRUCTURED_PROPERTIES_ASPECT_NAME,
+              STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME,
               FORM_INFO_ASPECT_NAME)
           .collect(Collectors.toSet());
 

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/template/structuredproperty/StructuredPropertyDefinitionTemplate.java

@@ -0,0 +1,86 @@
+package com.linkedin.metadata.aspect.patch.template.structuredproperty;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.linkedin.common.UrnArray;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.metadata.aspect.patch.template.CompoundKeyTemplate;
+import com.linkedin.structured.StructuredPropertyDefinition;
+import java.util.Collections;
+import javax.annotation.Nonnull;
+
+public class StructuredPropertyDefinitionTemplate
+    extends CompoundKeyTemplate<StructuredPropertyDefinition> {
+
+  private static final String ENTITY_TYPES_FIELD_NAME = "entityTypes";
+  private static final String ALLOWED_VALUES_FIELD_NAME = "allowedValues";
+  private static final String VALUE_FIELD_NAME = "value";
+  private static final String UNIT_SEPARATOR_DELIMITER = "␟";
+
+  @Override
+  public StructuredPropertyDefinition getSubtype(RecordTemplate recordTemplate)
+      throws ClassCastException {
+    if (recordTemplate instanceof StructuredPropertyDefinition) {
+      return (StructuredPropertyDefinition) recordTemplate;
+    }
+    throw new ClassCastException("Unable to cast RecordTemplate to StructuredPropertyDefinition");
+  }
+
+  @Override
+  public Class<StructuredPropertyDefinition> getTemplateType() {
+    return StructuredPropertyDefinition.class;
+  }
+
+  @Nonnull
+  @Override
+  public StructuredPropertyDefinition getDefault() {
+    StructuredPropertyDefinition definition = new StructuredPropertyDefinition();
+    definition.setQualifiedName("");
+    definition.setValueType(UrnUtils.getUrn("urn:li:dataType:datahub.string"));
+    definition.setEntityTypes(new UrnArray());
+
+    return definition;
+  }
+
+  @Nonnull
+  @Override
+  public JsonNode transformFields(JsonNode baseNode) {
+    JsonNode transformedNode =
+        arrayFieldToMap(baseNode, ENTITY_TYPES_FIELD_NAME, Collections.emptyList());
+
+    if (transformedNode.get(ALLOWED_VALUES_FIELD_NAME) == null) {
+      return transformedNode;
+    }
+
+    // allowedValues has a nested key - value.string or value.number depending on type. Mapping
+    // needs this nested key
+    JsonNode allowedValues = transformedNode.get(ALLOWED_VALUES_FIELD_NAME);
+    if (((ArrayNode) allowedValues).size() > 0) {
+      JsonNode allowedValue = ((ArrayNode) allowedValues).get(0);
+      JsonNode value = allowedValue.get(VALUE_FIELD_NAME);
+      String secondaryKeyName = value.fieldNames().next();
+      return arrayFieldToMap(
+          transformedNode,
+          ALLOWED_VALUES_FIELD_NAME,
+          Collections.singletonList(
+              VALUE_FIELD_NAME + UNIT_SEPARATOR_DELIMITER + secondaryKeyName));
+    }
+
+    return arrayFieldToMap(
+        transformedNode, ALLOWED_VALUES_FIELD_NAME, Collections.singletonList(VALUE_FIELD_NAME));
+  }
+
+  @Nonnull
+  @Override
+  public JsonNode rebaseFields(JsonNode patched) {
+    JsonNode patchedNode =
+        transformedMapToArray(patched, ENTITY_TYPES_FIELD_NAME, Collections.emptyList());
+
+    if (patchedNode.get(ALLOWED_VALUES_FIELD_NAME) == null) {
+      return patchedNode;
+    }
+    return transformedMapToArray(
+        patchedNode, ALLOWED_VALUES_FIELD_NAME, Collections.singletonList(VALUE_FIELD_NAME));
+  }
+}

entity-registry/src/main/java/com/linkedin/metadata/models/registry/SnapshotEntityRegistry.java

@@ -21,6 +21,7 @@
 import com.linkedin.metadata.aspect.patch.template.dataset.EditableSchemaMetadataTemplate;
 import com.linkedin.metadata.aspect.patch.template.dataset.UpstreamLineageTemplate;
 import com.linkedin.metadata.aspect.patch.template.form.FormInfoTemplate;
+import com.linkedin.metadata.aspect.patch.template.structuredproperty.StructuredPropertyDefinitionTemplate;
 import com.linkedin.metadata.models.AspectSpec;
 import com.linkedin.metadata.models.DefaultEntitySpec;
 import com.linkedin.metadata.models.EntitySpec;
@@ -88,6 +89,8 @@ private AspectTemplateEngine populateTemplateEngine(Map<String, AspectSpec> aspe
     aspectSpecTemplateMap.put(DATA_JOB_INPUT_OUTPUT_ASPECT_NAME, new DataJobInputOutputTemplate());
     aspectSpecTemplateMap.put(
         STRUCTURED_PROPERTIES_ASPECT_NAME, new StructuredPropertiesTemplate());
+    aspectSpecTemplateMap.put(
+        STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME, new StructuredPropertyDefinitionTemplate());
     aspectSpecTemplateMap.put(FORM_INFO_ASPECT_NAME, new FormInfoTemplate());
     return new AspectTemplateEngine(aspectSpecTemplateMap);
   }

metadata-ingestion/examples/library/dataset_add_structured_properties.py

@@ -0,0 +1,24 @@
+import logging
+
+from datahub.emitter.mce_builder import make_dataset_urn
+from datahub.emitter.rest_emitter import DataHubRestEmitter
+from datahub.specific.dataset import DatasetPatchBuilder
+
+log = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Create rest emitter
+rest_emitter = DataHubRestEmitter(gms_server="http://localhost:8080")
+
+dataset_urn = make_dataset_urn(platform="hive", name="fct_users_created", env="PROD")
+
+
+for patch_mcp in (
+    DatasetPatchBuilder(dataset_urn)
+    .add_structured_property("io.acryl.dataManagement.replicationSLA", 12)
+    .build()
+):
+    rest_emitter.emit(patch_mcp)
+
+
+log.info(f"Added cluster_name, retention_time properties to dataset {dataset_urn}")

metadata-ingestion/examples/library/dataset_remove_structured_properties.py

@@ -0,0 +1,24 @@
+import logging
+
+from datahub.emitter.mce_builder import make_dataset_urn
+from datahub.emitter.rest_emitter import DataHubRestEmitter
+from datahub.specific.dataset import DatasetPatchBuilder
+
+log = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Create rest emitter
+rest_emitter = DataHubRestEmitter(gms_server="http://localhost:8080")
+
+dataset_urn = make_dataset_urn(platform="hive", name="fct_users_created", env="PROD")
+
+
+for patch_mcp in (
+    DatasetPatchBuilder(dataset_urn)
+    .remove_structured_property("io.acryl.dataManagement.replicationSLA")
+    .build()
+):
+    rest_emitter.emit(patch_mcp)
+
+
+log.info(f"Added cluster_name, retention_time properties to dataset {dataset_urn}")

metadata-ingestion/examples/library/dataset_update_structured_properties.py

@@ -0,0 +1,24 @@
+import logging
+
+from datahub.emitter.mce_builder import make_dataset_urn
+from datahub.emitter.rest_emitter import DataHubRestEmitter
+from datahub.specific.dataset import DatasetPatchBuilder
+
+log = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Create rest emitter
+rest_emitter = DataHubRestEmitter(gms_server="http://localhost:8080")
+
+dataset_urn = make_dataset_urn(platform="hive", name="fct_users_created", env="PROD")
+
+
+for patch_mcp in (
+    DatasetPatchBuilder(dataset_urn)
+    .set_structured_property("io.acryl.dataManagement.replicationSLA", 120)
+    .build()
+):
+    rest_emitter.emit(patch_mcp)
+
+
+log.info(f"Added cluster_name, retention_time properties to dataset {dataset_urn}")

metadata-ingestion/examples/structured_properties/create_structured_property.py

@@ -0,0 +1,98 @@
+import logging
+
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.rest_emitter import DatahubRestEmitter
+
+# Imports for metadata model classes
+from datahub.metadata.schema_classes import (
+    PropertyValueClass,
+    StructuredPropertyDefinitionClass,
+)
+from datahub.metadata.urns import StructuredPropertyUrn
+
+log = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Create rest emitter
+rest_emitter = DatahubRestEmitter(gms_server="http://localhost:8080")
+
+# first, let's make an open ended structured property that allows one text value
+text_property_urn = StructuredPropertyUrn("openTextProperty")
+text_property_definition = StructuredPropertyDefinitionClass(
+    qualifiedName="io.acryl.openTextProperty",
+    displayName="Open Text Property",
+    valueType="urn:li:dataType:datahub.string",
+    cardinality="SINGLE",
+    entityTypes=[
+        "urn:li:entityType:datahub.dataset",
+        "urn:li:entityType:datahub.container",
+    ],
+    description="This structured property allows a signle open ended response as a value",
+    immutable=False,
+)
+
+event_prop_1: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
+    entityUrn=str(text_property_urn),
+    aspect=text_property_definition,
+)
+rest_emitter.emit(event_prop_1)
+
+# next, let's make a property that allows for multiple datahub entity urns as values
+# This example property could be used to reference other users or groups in datahub
+urn_property_urn = StructuredPropertyUrn("dataSteward")
+urn_property_definition = StructuredPropertyDefinitionClass(
+    qualifiedName="io.acryl.dataManagement.dataSteward",
+    displayName="Data Steward",
+    valueType="urn:li:dataType:datahub.urn",
+    cardinality="MULTIPLE",
+    entityTypes=["urn:li:entityType:datahub.dataset"],
+    description="The data stewards of this asset are in charge of ensuring data cleanliness and governance",
+    immutable=True,
+    typeQualifier={
+        "allowedTypes": [
+            "urn:li:entityType:datahub.corpuser",
+            "urn:li:entityType:datahub.corpGroup",
+        ]
+    },  # this line ensures only user or group urns can be assigned as values
+)
+
+event_prop_2: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
+    entityUrn=str(urn_property_urn),
+    aspect=urn_property_definition,
+)
+rest_emitter.emit(event_prop_2)
+
+# finally, let's make a single select number property with a few allowed options
+number_property_urn = StructuredPropertyUrn("replicationSLA")
+number_property_definition = StructuredPropertyDefinitionClass(
+    qualifiedName="io.acryl.dataManagement.replicationSLA",
+    displayName="Retention Time",
+    valueType="urn:li:dataType:datahub.number",
+    cardinality="SINGLE",
+    entityTypes=[
+        "urn:li:entityType:datahub.dataset",
+        "urn:li:entityType:datahub.dataFlow",
+    ],
+    description="SLA for how long data can be delayed before replicating to the destination cluster",
+    immutable=False,
+    allowedValues=[
+        PropertyValueClass(
+            value=30,
+            description="30 days, usually reserved for datasets that are ephemeral and contain pii",
+        ),
+        PropertyValueClass(
+            value=90,
+            description="Use this for datasets that drive monthly reporting but contain pii",
+        ),
+        PropertyValueClass(
+            value=365,
+            description="Use this for non-sensitive data that can be retained for longer",
+        ),
+    ],
+)
+
+event_prop_3: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
+    entityUrn=str(number_property_urn),
+    aspect=number_property_definition,
+)
+rest_emitter.emit(event_prop_3)

metadata-ingestion/examples/structured_properties/update_structured_property.py

@@ -0,0 +1,43 @@
+import logging
+from typing import Union
+
+from datahub.configuration.kafka import KafkaProducerConnectionConfig
+from datahub.emitter.kafka_emitter import DatahubKafkaEmitter, KafkaEmitterConfig
+from datahub.emitter.rest_emitter import DataHubRestEmitter
+from datahub.metadata.urns import StructuredPropertyUrn
+from datahub.specific.structured_property import StructuredPropertyPatchBuilder
+
+log = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+# Get an emitter, either REST or Kafka, this example shows you both
+def get_emitter() -> Union[DataHubRestEmitter, DatahubKafkaEmitter]:
+    USE_REST_EMITTER = True
+    if USE_REST_EMITTER:
+        gms_endpoint = "http://localhost:8080"
+        return DataHubRestEmitter(gms_server=gms_endpoint)
+    else:
+        kafka_server = "localhost:9092"
+        schema_registry_url = "http://localhost:8081"
+        return DatahubKafkaEmitter(
+            config=KafkaEmitterConfig(
+                connection=KafkaProducerConnectionConfig(
+                    bootstrap=kafka_server, schema_registry_url=schema_registry_url
+                )
+            )
+        )
+
+
+# input your unique structured property ID
+property_urn = StructuredPropertyUrn("dataSteward")
+
+with get_emitter() as emitter:
+    for patch_mcp in (
+        StructuredPropertyPatchBuilder(str(property_urn))
+        .set_display_name("test display name")
+        .set_cardinality("MULTIPLE")
+        .add_entity_type("urn:li:entityType:datahub.dataJob")
+        .build()
+    ):
+        emitter.emit(patch_mcp)