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 index 8d7c6b2a80..30d812b1ef 100644 --- 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 @@ -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 private val header: List = CSVXformer.getHeader(filename, fieldSeparator) private val tagIdx: Int = header.indexOf(TAG_NAME) + private val extraColumns: List = 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, + ): 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, + ) { + // 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, header: List, diff --git a/samples/packages/asset-import/src/test/kotlin/com/atlan/pkg/aim/TagExtraAttributesTest.kt b/samples/packages/asset-import/src/test/kotlin/com/atlan/pkg/aim/TagExtraAttributesTest.kt new file mode 100644 index 0000000000..60e8df72a5 --- /dev/null +++ b/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() + } +} diff --git a/samples/packages/asset-import/src/test/resources/tags_with_extra_attrs.csv b/samples/packages/asset-import/src/test/resources/tags_with_extra_attrs.csv new file mode 100644 index 0000000000..6d9abcf263 --- /dev/null +++ b/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""}"