Adds generic SourceTag handling

Signed-off-by: Christopher Grote <[email protected]>
atlanhq · Feb 13, 2025 · f1a2e8d · f1a2e8d
1 parent 2a38572
commit f1a2e8d
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 25 deletions.
diff --git a/samples/packages/asset-import/src/main/kotlin/com/atlan/pkg/aim/AtlanTagImporter.kt b/samples/packages/asset-import/src/main/kotlin/com/atlan/pkg/aim/AtlanTagImporter.kt
@@ -182,27 +182,28 @@ class AtlanTagImporter(
 
     private fun idempotentTagAsset(tag: TagDetails): Asset? {
         return if (tag.sourceSynced) {
-            val assetBuilder = when (tag.connectorType) {
-                "snowflake" -> {
-                    // TODO: actually use a schema for Snowflake, not just a connection
-                    SnowflakeTag.creator(
-                        tag.name,
-                        tag.connectionQualifiedName,
-                        tag.name,
-                        tag.tagIdInSource,
-                        tag.allowedValues,
-                    )
-                }
-                else -> {
-                    SourceTag.creator(
-                        tag.name,
-                        tag.connectionQualifiedName,
-                        tag.name,
-                        tag.tagIdInSource,
-                        tag.allowedValues,
-                    )
+            val assetBuilder =
+                when (tag.connectorType) {
+                    "snowflake" -> {
+                        // TODO: actually use a schema for Snowflake, not just a connection
+                        SnowflakeTag.creator(
+                            tag.name,
+                            tag.connectionQualifiedName,
+                            tag.name,
+                            tag.tagIdInSource,
+                            tag.allowedValues,
+                        )
+                    }
+                    else -> {
+                        SourceTag.creator(
+                            tag.name,
+                            tag.connectionQualifiedName,
+                            tag.name,
+                            tag.tagIdInSource,
+                            tag.allowedValues,
+                        )
+                    }
                 }
-            }
             assetBuilder.build()
         } else {
             null

diff --git a/sdk/src/main/java/com/atlan/model/assets/SnowflakeTag.java b/sdk/src/main/java/com/atlan/model/assets/SnowflakeTag.java
@@ -553,7 +553,10 @@ public static boolean restore(AtlanClient client, String qualifiedName) throws A
                 .connectionQualifiedName(connectionQualifiedName)
                 .mappedAtlanTagName(mappedAtlanTagName)
                 .tagId(snowflakeTagId)
-                .tagAttribute(SourceTagAttribute.of("allowedValues", allowedValuesString, null))
+                .tagAttribute(SourceTagAttribute.builder()
+                        .tagAttributeKey("allowedValues")
+                        .tagAttributeValue(allowedValuesString)
+                        .build())
                 .tagAllowedValues(allowedValues);
     }
 

diff --git a/sdk/src/main/java/com/atlan/model/assets/SourceTag.java b/sdk/src/main/java/com/atlan/model/assets/SourceTag.java
@@ -378,7 +378,10 @@ public static boolean restore(AtlanClient client, String qualifiedName) throws A
                 .connectionQualifiedName(connectionQualifiedName)
                 .mappedAtlanTagName(mappedAtlanTagName)
                 .tagId(sourceTagId)
-                .tagAttribute(SourceTagAttribute.of("allowedValues", allowedValuesString, null))
+                .tagAttribute(SourceTagAttribute.builder()
+                        .tagAttributeKey("allowedValues")
+                        .tagAttributeValue(allowedValuesString)
+                        .build())
                 .tagAllowedValues(allowedValues);
     }
 
@@ -400,7 +403,7 @@ public static boolean restore(AtlanClient client, String qualifiedName) throws A
      * Generate a unique SourceTag name.
      *
      * @param name of the SourceTag
-     * @param connectionQualifiedName unique name of the connection in which this SourceTag exists
+     * @param connectionQualifiedName unique name of the schema in which this SourceTag exists
      * @return a unique name for the SourceTag
      */
     public static String generateQualifiedName(String name, String connectionQualifiedName) {

diff --git a/sdk/src/main/resources/templates/SnowflakeTag.ftl b/sdk/src/main/resources/templates/SnowflakeTag.ftl
@@ -95,7 +95,7 @@
             .connectionQualifiedName(connectionQualifiedName)
             .mappedAtlanTagName(mappedAtlanTagName)
             .tagId(snowflakeTagId)
-            .tagAttribute(SourceTagAttribute.of("allowedValues", allowedValuesString, null))
+            .tagAttribute(SourceTagAttribute.builder().tagAttributeKey("allowedValues").tagAttributeValue(allowedValuesString).build())
             .tagAllowedValues(allowedValues);
     }
 

diff --git a/sdk/src/main/resources/templates/SourceTag.ftl b/sdk/src/main/resources/templates/SourceTag.ftl
@@ -31,7 +31,7 @@
             .connectionQualifiedName(connectionQualifiedName)
             .mappedAtlanTagName(mappedAtlanTagName)
             .tagId(sourceTagId)
-            .tagAttribute(SourceTagAttribute.of("allowedValues", allowedValuesString, null))
+            .tagAttribute(SourceTagAttribute.builder().tagAttributeKey("allowedValues").tagAttributeValue(allowedValuesString).build())
             .tagAllowedValues(allowedValues);
     }