Merge pull request #1323 from atlanhq/APP-5462

Adds generic SourceTag and SnowflakeTag handling
atlanhq · Feb 13, 2025 · 62e054e · 62e054e
2 parents ae8f0dc + 71bcfd3
commit 62e054e
Show file tree

Hide file tree

Showing 582 changed files with 8,141 additions and 54 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
@@ -6,7 +6,7 @@ import AssetImportCfg
 import com.atlan.exception.NotFoundException
 import com.atlan.model.assets.Asset
 import com.atlan.model.assets.SnowflakeTag
-import com.atlan.model.enums.AtlanConnectorType
+import com.atlan.model.assets.SourceTag
 import com.atlan.model.enums.AtlanEnum
 import com.atlan.model.enums.AtlanIcon
 import com.atlan.model.enums.AtlanTagColor
@@ -26,7 +26,6 @@ import de.siegmar.fastcsv.reader.CsvRecord
 import mu.KLogger
 import java.io.IOException
 import java.nio.file.Paths
-import java.util.concurrent.ThreadLocalRandom
 import java.util.concurrent.atomic.AtomicLong
 import kotlin.streams.asSequence
 
@@ -104,6 +103,7 @@ class AtlanTagImporter(
                                 row.getOrElse(TAG_ICON) { "" },
                                 row.getOrElse(TAG_CONNECTION) { "" },
                                 row.getOrElse(TAG_CONNECTOR) { "" }.lowercase(),
+                                row.getOrElse(TAG_SRC_ID) { "" },
                                 CellXformer.parseDelimitedList(row.getOrElse(ALLOWED_VALUES) { "" }),
                             )
 
@@ -181,8 +181,34 @@ class AtlanTagImporter(
     }
 
     private fun idempotentTagAsset(tag: TagDetails): Asset? {
-        // TODO: replace with creator() call, and generic SourceTag — should be idempotent already (?)
         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,
+                        )
+                    }
+                }
+            assetBuilder.build()
+        } else {
+            null
+        }
+        /*return if (tag.sourceSynced) {
             // Source tag asset
             val assetBuilder =
                 SnowflakeTag
@@ -200,7 +226,7 @@ class AtlanTagImporter(
             assetBuilder.build()
         } else {
             null
-        }
+        }*/
     }
 
     companion object {
@@ -209,17 +235,8 @@ class AtlanTagImporter(
         const val TAG_ICON = "Icon"
         const val TAG_CONNECTION = "Synced connection name"
         const val TAG_CONNECTOR = "Synced connector type"
+        const val TAG_SRC_ID = "Tag ID in source"
         const val ALLOWED_VALUES = "Allowed values for tag"
-        val HEADER =
-            listOf(
-                TAG_NAME,
-                TAG_COLOR,
-                TAG_ICON,
-                TAG_CONNECTION,
-                TAG_CONNECTOR,
-                ALLOWED_VALUES,
-            )
-        const val NO_CONNECTION_QN = "NO_CONNECTION_FOUND"
 
         fun getTagName(
             row: List<String>,
@@ -251,6 +268,7 @@ class AtlanTagImporter(
         val icon: String,
         val connectionName: String,
         val connectorType: String,
+        val tagIdInSource: String,
         val allowedValues: List<String>,
     ) {
         val connectionQualifiedName: String

diff --git a/samples/packages/asset-import/src/test/kotlin/com/atlan/pkg/aim/TagManagementTest.kt b/samples/packages/asset-import/src/test/kotlin/com/atlan/pkg/aim/TagManagementTest.kt
@@ -24,7 +24,7 @@ class TagManagementTest : PackageTest("tm") {
     override val logger = Utils.getLogger(this.javaClass.name)
 
     private val c1 = makeUnique("c1")
-    private val c1Type = AtlanConnectorType.SNOWFLAKE // TODO: swap to BigID once generic SourceTag is in place
+    private val c1Type = AtlanConnectorType.BIGID
     private val t1 = makeUnique("t1")
     private val t2 = makeUnique("t2")
 

diff --git a/samples/packages/asset-import/src/test/resources/tags.csv b/samples/packages/asset-import/src/test/resources/tags.csv
@@ -1,5 +1,5 @@
-Atlan tag name,Color,Icon,Synced connection name,Synced connector type,Allowed values for tag
+Atlan tag name,Color,Icon,Synced connection name,Synced connector type,Allowed values for tag,Tag ID in source
 {{TAG1}},Green,PhRecycle,{{CONNECTION}},{{CTYPE}},"X
 Y
-Z"
-{{TAG2}},,https://assets.atlan.com/assets/announcement-issue-png.png,,,
+Z",42
+{{TAG2}},,https://assets.atlan.com/assets/announcement-issue-png.png,,,,
diff --git a/sdk/src/main/java/com/atlan/generators/GeneratorConfig.java b/sdk/src/main/java/com/atlan/generators/GeneratorConfig.java
@@ -83,7 +83,8 @@ public class GeneratorConfig {
             Map.entry("exceptionsForBusinessPolicy", "exceptionForBusinessPolicy"),
             Map.entry("supersetChartFormData", "putSupersetChartFormData"),
             Map.entry("columnHierarchy", "putColumnHierarchy"),
-            Map.entry("bigqueryTagHierarchy", "putBigqueryTagHierarchy"));
+            Map.entry("bigqueryTagHierarchy", "putBigqueryTagHierarchy"),
+            Map.entry("applications", "aiApplication"));
 
     private static final Map<String, String> DEFAULT_ATTRIBUTE_RENAMES = Map.ofEntries(
             Map.entry("connectorName", "connectorType"),

diff --git a/sdk/src/main/java/com/atlan/model/assets/ADLSAccount.java b/sdk/src/main/java/com/atlan/model/assets/ADLSAccount.java
@@ -65,6 +65,10 @@ public class ADLSAccount extends Asset
     @Attribute
     ADLSStorageKind adlsAccountKind;
 
+    /** Name of the account for this ADLS asset. */
+    @Attribute
+    String adlsAccountName;
+
     /** Performance of this account. */
     @Attribute
     ADLSPerformance adlsAccountPerformance;

diff --git a/sdk/src/main/java/com/atlan/model/assets/ADLSContainer.java b/sdk/src/main/java/com/atlan/model/assets/ADLSContainer.java
@@ -56,6 +56,10 @@ public class ADLSContainer extends Asset
     @Attribute
     IADLSAccount adlsAccount;
 
+    /** Name of the account for this ADLS asset. */
+    @Attribute
+    String adlsAccountName;
+
     /** Unique name of the account for this ADLS asset. */
     @Attribute
     String adlsAccountQualifiedName;

diff --git a/sdk/src/main/java/com/atlan/model/assets/ADLSObject.java b/sdk/src/main/java/com/atlan/model/assets/ADLSObject.java
@@ -56,6 +56,10 @@ public class ADLSObject extends Asset
     @Builder.Default
     String typeName = TYPE_NAME;
 
+    /** Name of the account for this ADLS asset. */
+    @Attribute
+    String adlsAccountName;
+
     /** Unique name of the account for this ADLS asset. */
     @Attribute
     String adlsAccountQualifiedName;
@@ -68,6 +72,10 @@ public class ADLSObject extends Asset
     @Attribute
     IADLSContainer adlsContainer;
 
+    /** Name of the container this object exists within. */
+    @Attribute
+    String adlsContainerName;
+
     /** Unique name of the container this object exists within. */
     @Attribute
     String adlsContainerQualifiedName;