Merge pull request #650 from atlanhq/DVX-481

Adds domains to custom metadata extension package
atlanhq · May 23, 2024 · cc7462e · cc7462e
2 parents b3d4f6d + b981079
commit cc7462e
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 3 deletions.
diff --git a/samples/packages/custom-metadata-extender/src/main/kotlin/CustomMetadataExtender.kt b/samples/packages/custom-metadata-extender/src/main/kotlin/CustomMetadataExtender.kt
@@ -3,8 +3,10 @@
 import com.atlan.Atlan
 import com.atlan.exception.AtlanException
 import com.atlan.exception.NotFoundException
+import com.atlan.model.assets.DataDomain
 import com.atlan.model.assets.Glossary
 import com.atlan.model.typedefs.AttributeDef
+import com.atlan.model.typedefs.AttributeDefOptions
 import com.atlan.pkg.Utils
 import mu.KotlinLogging
 import kotlin.system.exitProcess
@@ -22,6 +24,8 @@ object CustomMetadataExtender {
         val cmName = Utils.getOrDefault(config.customMetadata, "")
         val connectionQNs = Utils.getOrDefault(config.connectionQualifiedName, listOf())
         val glossaryNames = Utils.getAsList(config.glossaries)
+        val domains = Utils.getOrDefault(config.domains, "ALL")
+        val domainName = Utils.getOrDefault(config.domainsSpecific, "")
 
         if (cmName.isBlank()) {
             logger.error { "Missing required parameter - you must specify the name of the custom metadata to extend." }
@@ -33,7 +37,7 @@ object CustomMetadataExtender {
             exitProcess(4)
         }
 
-        extendCM(cmName, connectionQNs, glossaryNames)
+        extendCM(cmName, connectionQNs, glossaryNames, domains, domainName)
     }
 
     /**
@@ -64,11 +68,40 @@ object CustomMetadataExtender {
      * @param cmName human-readable name of the custom metadata to extend
      * @param connectionQNs list of qualifiedNames of connections to add to the custom metadata
      * @param glossaryNames list of names of glossaries to add to the custom metadata
+     * @param domains whether to extend to all, some, or no domains
+     * @param domainName name of a single domain to extend to (if domains is some)
      */
-    fun extendCM(cmName: String, connectionQNs: List<String>, glossaryNames: List<String>) {
+    fun extendCM(
+        cmName: String,
+        connectionQNs: List<String>,
+        glossaryNames: List<String>,
+        domains: String,
+        domainName: String,
+    ) {
         logger.info { "Extending custom metadata $cmName with connections: $connectionQNs" }
         val glossaryQNs = getGlossaryQualifiedNames(glossaryNames)
         logger.info { "Extending custom metadata $cmName with glossaries: $glossaryQNs" }
+        val domainQNs = mutableSetOf<String>()
+        when (domains) {
+            "ALL" -> {
+                logger.info { "Extending custom metadata to all domains." }
+                domainQNs.addAll(AttributeDefOptions.ALL_DOMAINS)
+            }
+            "SOME" -> {
+                try {
+                    val found = DataDomain.findByName(domainName)
+                    if (found.size > 1) {
+                        logger.warn { "Found multiple domains with the name $domainName, taking only the first" }
+                    }
+                    domainQNs.add(found[0].qualifiedName)
+                } catch (e: NotFoundException) {
+                    logger.warn { "Unable to find domain with name $domainName -- skipping." }
+                } catch (e: AtlanException) {
+                    logger.error(e) { "Error attempting to lookup domain with name $domainName" }
+                }
+            }
+            else -> logger.info { "Not extending custom metadata to any additional domains." }
+        }
         val cm = Atlan.getDefaultClient().customMetadataCache.getCustomMetadataDef(cmName)
         if (cm == null) {
             logger.error { "Unable to find custom metadata with name: $cmName" }
@@ -78,6 +111,7 @@ object CustomMetadataExtender {
                 val options = attr.options.toBuilder()
                     .applicableConnections(connectionQNs)
                     .applicableGlossaries(glossaryQNs)
+                    .applicableDomains(domainQNs)
                     .build()
                 attrs.add(
                     attr.toBuilder()

diff --git a/samples/packages/custom-metadata-extender/src/main/kotlin/CustomMetadataExtenderCfg.kt b/samples/packages/custom-metadata-extender/src/main/kotlin/CustomMetadataExtenderCfg.kt
@@ -19,4 +19,6 @@ data class CustomMetadataExtenderCfg(
     @JsonSerialize(using = WidgetSerde.MultiSelectSerializer::class)
     @JsonProperty("connection_qualified_name") val connectionQualifiedName: List<String>? = null,
     @JsonProperty("glossaries") val glossaries: String? = null,
+    @JsonProperty("domains") val domains: String? = null,
+    @JsonProperty("domains_specific") val domainsSpecific: String? = null,
 ) : CustomConfig()
diff --git a/samples/packages/custom-metadata-extender/src/main/resources/package.pkl b/samples/packages/custom-metadata-extender/src/main/resources/package.pkl
@@ -56,7 +56,29 @@ uiConfig {
           width = 8
           placeholderText = "First Glossary,Second Glossary"
         }
+        ["domains"] = new Radio {
+          title = "Domains"
+          possibleValues {
+            ["ALL"] = "All domains"
+            ["SOME"] = "Specific domains"
+            ["NONE"] = "No domains"
+          }
+          default = "ALL"
+        }
+        ["domains_specific"] = new TextInput {
+          title = "Specific domain"
+          required = false
+          helpText = "Enter the name of the domain whose assets should be able to use the custom metadata."
+          width = 8
+          placeholderText = "Finance"
+        }
       }
     }
   }
+  rules {
+    new UIRule {
+      whenInputs { ["domains"] = "SOME" }
+      required { "domains_specific" }
+    }
+  }
 }
diff --git a/sdk/src/main/java/com/atlan/model/typedefs/AttributeDefOptions.java b/sdk/src/main/java/com/atlan/model/typedefs/AttributeDefOptions.java
@@ -227,6 +227,7 @@ public class AttributeDefOptions extends AtlanObject {
             Set.of(Glossary.TYPE_NAME, GlossaryTerm.TYPE_NAME, GlossaryCategory.TYPE_NAME);
     public static final Set<String> ALL_DOMAIN_TYPES = Set.of(DataDomain.TYPE_NAME, DataProduct.TYPE_NAME);
     public static final Set<String> ALL_OTHER_TYPES = Set.of(File.TYPE_NAME);
+    public static final Set<String> ALL_DOMAINS = Set.of("*/super");
 
     /**
      * Instantiate a new set of attribute options from the provided parameters.
@@ -262,7 +263,7 @@ public static AttributeDefOptions of(
                     .applicableAssetTypes(ALL_ASSET_TYPES)
                     .applicableGlossaries(Glossary.getAllQualifiedNames())
                     .applicableGlossaryTypes(ALL_GLOSSARY_TYPES)
-                    .applicableDomains(Set.of("*/super"))
+                    .applicableDomains(ALL_DOMAINS)
                     .applicableDomainTypes(ALL_DOMAIN_TYPES)
                     .applicableOtherAssetTypes(ALL_OTHER_TYPES);
         }