Adds support for additional attributes on tags in asset import

samples/packages/asset-import/src/main/kotlin/com/atlan/pkg/aim/AtlanTagImporter.kt

@@ -3,6 +3,7 @@
 package com.atlan.pkg.aim
 
 import AssetImportCfg
+import com.atlan.cache.ReflectionCache
 import com.atlan.exception.NotFoundException
 import com.atlan.model.assets.Asset
 import com.atlan.model.assets.DatabricksUnityCatalogTag
@@ -17,6 +18,7 @@ import com.atlan.model.typedefs.AtlanTagOptions
 import com.atlan.model.typedefs.TypeDef
 import com.atlan.pkg.PackageContext
 import com.atlan.pkg.Utils
+import com.atlan.pkg.serde.FieldSerde
 import com.atlan.pkg.serde.cell.CellXformer
 import com.atlan.pkg.serde.cell.ConnectionXformer
 import com.atlan.pkg.serde.cell.EnumXformer
@@ -54,6 +56,7 @@ class AtlanTagImporter(
     private val counter: CsvReader<CsvRecord>
     private val header: List<String> = CSVXformer.getHeader(filename, fieldSeparator)
     private val tagIdx: Int = header.indexOf(TAG_NAME)
+    private val extraColumns: List<String> = header.filter { !KNOWN_COLUMNS.contains(it) && it.isNotBlank() }
 
     init {
         val missingColumns = validateHeader(header)
@@ -140,7 +143,7 @@ class AtlanTagImporter(
                         count.getAndIncrement()
 
                         // Manage the asset, for any source-synced tags
-                        val asset = idempotentTagAsset(details)
+                        val asset = idempotentTagAsset(details, row)
                         if (asset != null) {
                             assets.add(asset)
                         }
@@ -212,7 +215,10 @@ class AtlanTagImporter(
             AtlanTagOptions.of(AtlanTagColor.GRAY, sourceSynced)
         }
 
-    private fun idempotentTagAsset(tag: TagDetails): Asset? =
+    private fun idempotentTagAsset(
+        tag: TagDetails,
+        row: Map<String, String>,
+    ): Asset? =
         if (tag.sourceSynced) {
             val assetBuilder =
                 when (tag.connectorType) {
@@ -262,11 +268,42 @@ class AtlanTagImporter(
                         )
                     }
                 }
+            // Apply any extra columns as attributes on the SourceTag asset
+            applyExtraAttributes(assetBuilder, row)
             assetBuilder.build()
         } else {
             null
         }
 
+    /**
+     * Apply extra columns from the CSV row as attributes on the asset builder using reflection.
+     * This allows any SourceTag attribute to be set via CSV columns without explicit handling.
+     *
+     * @param builder the asset builder to apply attributes to
+     * @param row the CSV row data as a map of column name to value
+     */
+    private fun applyExtraAttributes(
+        builder: Asset.AssetBuilder<*, *>,
+        row: Map<String, String>,
+    ) {
+        // Get the builder class from the actual builder instance to handle SnowflakeTag, DbtTag, etc.
+        val builderClass = builder.javaClass
+        for (columnName in extraColumns) {
+            val rawValue = CSVXformer.trimWhitespace(row.getOrElse(columnName) { "" })
+            if (rawValue.isNotBlank()) {
+                val setter = ReflectionCache.getSetter(builderClass, columnName)
+                if (setter != null) {
+                    val value = FieldSerde.getValueFromCell(ctx, rawValue, setter, logger)
+                    if (value != null) {
+                        ReflectionCache.setValue(builder, columnName, value)
+                    }
+                } else {
+                    logger.warn { "Unknown column '$columnName' in tags CSV - no matching attribute found on SourceTag, skipping." }
+                }
+            }
+        }
+    }
+
     companion object {
         const val TAG_NAME = "Atlan tag name"
         const val TAG_COLOR = "Color"
@@ -279,6 +316,23 @@ class AtlanTagImporter(
         const val DBT_PROJECT_ID = "Project ID (dbt)"
         const val SNOWFLAKE_PATH = "Schema path (Snowflake)"
 
+        /** Set of known columns that are handled explicitly by the importer. */
+        val KNOWN_COLUMNS =
+            setOf(
+                TAG_NAME,
+                TAG_COLOR,
+                TAG_ICON,
+                TAG_CONNECTION,
+                TAG_CONNECTOR,
+                TAG_SRC_ID,
+                ALLOWED_VALUES,
+                DBT_ACCOUNT_ID,
+                DBT_PROJECT_ID,
+                SNOWFLAKE_PATH,
+                "Description",
+                "description",
+            )
+
         fun getTagName(
             row: List<String>,
             header: List<String>,

samples/packages/asset-import/src/test/kotlin/com/atlan/pkg/aim/TagExtraAttributesTest.kt

@@ -0,0 +1,117 @@
+/* SPDX-License-Identifier: Apache-2.0
+   Copyright 2023 Atlan Pte. Ltd. */
+package com.atlan.pkg.aim
+
+import AssetImportCfg
+import com.atlan.model.assets.Connection
+import com.atlan.model.assets.SourceTag
+import com.atlan.model.enums.AtlanConnectorType
+import com.atlan.model.enums.AtlanIcon
+import com.atlan.model.enums.AtlanTagColor
+import com.atlan.model.enums.TagIconType
+import com.atlan.pkg.PackageTest
+import com.atlan.pkg.Utils
+import java.nio.file.Paths
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+
+/**
+ * Test import of Atlan tags with extra attributes via CSV columns.
+ * This tests the ability to set arbitrary SourceTag attributes (like tagCustomConfiguration)
+ * through CSV columns that are not explicitly handled by the importer.
+ */
+class TagExtraAttributesTest : PackageTest("tea") {
+    override val logger = Utils.getLogger(this.javaClass.name)
+
+    private val c1 = makeUnique("c1")
+    private val c1Type = AtlanConnectorType.BIGID
+    private val t1 = makeUnique("t1")
+
+    private lateinit var connection: Connection
+
+    private val tagsFile = "tags_with_extra_attrs.csv"
+
+    private val files =
+        listOf(
+            tagsFile,
+            "debug.log",
+        )
+
+    private fun prepFile() {
+        // Prepare a copy of the file with unique names for objects
+        val tagsIn = Paths.get("src", "test", "resources", tagsFile).toFile()
+        val tagsOut = Paths.get(testDirectory, tagsFile).toFile()
+        tagsIn.useLines { lines ->
+            lines.forEach { line ->
+                val revised =
+                    line
+                        .replace("{{TAG1}}", t1)
+                        .replace("{{CONNECTION}}", c1)
+                        .replace("{{CTYPE}}", c1Type.value)
+                tagsOut.appendText("$revised\n")
+            }
+        }
+    }
+
+    private fun createConnection(): Connection {
+        val conn = Connection.creator(client, c1, c1Type).build()
+        val response = conn.save(client).block()
+        return response.getResult(conn)
+    }
+
+    override fun setup() {
+        connection = createConnection()
+        prepFile()
+        runCustomPackage(
+            AssetImportCfg(
+                tagsFile = Paths.get(testDirectory, tagsFile).toString(),
+                tagsFailOnErrors = true,
+            ),
+            Importer::main,
+        )
+    }
+
+    override fun teardown() {
+        removeConnection(c1, c1Type)
+        removeTag(t1)
+    }
+
+    @Test
+    fun tagExists() {
+        val tag = client.atlanTagCache.getByName(t1)
+        assertNotNull(tag)
+        assertEquals(t1, tag.displayName)
+        assertEquals("Icon tag with extra config", tag.description)
+        assertEquals(AtlanTagColor.GREEN.value, tag.options.color)
+        assertEquals(AtlanIcon.RECYCLE, tag.options.iconName)
+        assertEquals(TagIconType.ICON, tag.options.iconType)
+    }
+
+    @Test
+    fun sourceTagHasExtraAttribute() {
+        val request =
+            SourceTag
+                .select(client)
+                .where(SourceTag.CONNECTION_QUALIFIED_NAME.eq(connection.qualifiedName))
+                .where(SourceTag.NAME.eq(t1))
+                .includeOnResults(SourceTag.TAG_CUSTOM_CONFIGURATION)
+                .toRequest()
+        val response = retrySearchUntil(request, 1)
+        val found = response.assets
+        assertEquals(1, found.size)
+        val sourceTag = found[0] as SourceTag
+        assertNotNull(sourceTag.tagCustomConfiguration)
+        assertEquals("{\"customKey\": \"customValue\"}", sourceTag.tagCustomConfiguration)
+    }
+
+    @Test
+    fun filesCreated() {
+        validateFilesExist(files)
+    }
+
+    @Test
+    fun errorFreeLog() {
+        validateErrorFreeLog()
+    }
+}

samples/packages/asset-import/src/test/resources/tags_with_extra_attrs.csv

@@ -0,0 +1,4 @@
+Atlan tag name,Color,Icon,Synced connection name,Synced connector type,Allowed values for tag,Tag ID in source,Account ID (dbt),Project ID (dbt),Schema path (Snowflake),Description,tagCustomConfiguration
+{{TAG1}},Green,PhRecycle,{{CONNECTION}},{{CTYPE}},"X
+Y
+Z",42,,,,Icon tag with extra config,"{""customKey"": ""customValue""}"