From fa6f939a483c370eedb9b461ba660c8e4f7bb3fb Mon Sep 17 00:00:00 2001 From: Vitaliano Vidaurre <122127166+vitaliano-atlan@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:55:06 -0600 Subject: [PATCH 1/2] fix(lake-formation-tag-sync): match views when syncing LF tags Lake Formation lists views as tables and the tag feed carries no table type, so every table-level row is written as Table. In update-only mode AssetBatch matches on type plus qualifiedName, so rows for existing views were skipped while the run still reported success. Turn on assetsTableViewAgnostic so a Table row that is not found as a Table matches a View, MaterializedView or SnowflakeDynamicTable. Extend LakeTagSynchronizerTest with a view and a view column. Refs CSA-614 Co-Authored-By: Claude Opus 5.5 --- .../atlan/pkg/lftag/LakeTagSynchronizer.kt | 3 ++ .../pkg/lftag/LakeTagSynchronizerTest.kt | 26 +++++++++++ .../test/resources/lftag_association_1.json | 46 +++++++++++++++++++ 3 files changed, 75 insertions(+) diff --git a/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt b/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt index bd2ddf16df..96329cedcc 100644 --- a/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt +++ b/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt @@ -90,6 +90,9 @@ object LakeTagSynchronizer { AssetImportCfg( assetsFile = csvFileName, assetsUpsertSemantic = "update", + // Lake Formation reports views as tables, and the tag feed carries no table type, + // so let a Table row match an existing View or MaterializedView of the same name. + assetsTableViewAgnostic = true, assetsConfig = "advanced", assetsFailOnErrors = ctx.config.getEffectiveValue( diff --git a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt index c852c8dd4e..817e41bd13 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt +++ b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt @@ -8,6 +8,7 @@ import com.atlan.model.assets.Connection import com.atlan.model.assets.Database import com.atlan.model.assets.Schema import com.atlan.model.assets.Table +import com.atlan.model.assets.View import com.atlan.model.enums.AtlanConnectorType import com.atlan.model.enums.AtlanCustomAttributePrimitiveType import com.atlan.model.typedefs.AttributeDef @@ -39,6 +40,8 @@ class LakeTagSynchronizerTest : PackageTest("lts") { private var schemaGuid = "" private var tableGuid = "" private var columnGuid = "" + private var viewGuid = "" + private var viewColumnGuid = "" private val cm1 = makeUnique("cm") private val enum1 = makeUnique("enum") @@ -71,10 +74,16 @@ class LakeTagSynchronizerTest : PackageTest("lts") { batch.add(tbl) val column = Column.creator("col1", tbl, 1).build() batch.add(column) + val vw = View.creator("vw1", sch).build() + batch.add(vw) + val vwColumn = Column.creator("vcol1", vw, 1).build() + batch.add(vwColumn) val response = batch.flush() schemaGuid = response.getResult(sch).guid tableGuid = response.getResult(tbl).guid columnGuid = response.getResult(column).guid + viewGuid = response.getResult(vw).guid + viewColumnGuid = response.getResult(vwColumn).guid } } @@ -177,6 +186,23 @@ class LakeTagSynchronizerTest : PackageTest("lts") { assertEquals(PUBLIC, attribute1) } + @Test + fun validateViewTagged() { + val view = View.get(client, viewGuid) + val attribute1 = view.getCustomMetadata(cm1, attr1) + assertEquals(PUBLIC, attribute1) + val attribute3 = view.getCustomMetadata(cm1, attr3) + assertEquals(FULL_HISTORY, attribute3) + } + + @Test + fun validateViewColumnTagged() { + val column = Column.get(client, viewColumnGuid) + assertEquals("vcol1", column.name) + val attribute1 = column.getCustomMetadata(cm1, attr1) + assertEquals(PUBLIC, attribute1) + } + @Test fun validateFilesCreated() { validateFilesExist( diff --git a/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json b/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json index af5e0b56b9..e1820b98c2 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json +++ b/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json @@ -82,6 +82,52 @@ "DatabaseName": "dev_sch", "Name": "sch.tbl1" } + }, + { + "Table": { + "CatalogId": "614518280298", + "DatabaseName": "dev_sch", + "Name": "sch.vw1" + }, + "LFTagOnDatabase": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + } + ], + "LFTagsOnTable": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + }, + { + "CatalogId": "614518280298", + "TagKey": "data_load_method", + "TagValues": [ + "fullhistory" + ] + } + ], + "LFTagsOnColumns": [ + { + "Name": "vcol1", + "LFTags": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + } + ] + } + ] } ] } From 413eb37ce4f8c394d33ccb9d7da6783df29006b5 Mon Sep 17 00:00:00 2001 From: Vitaliano Vidaurre <122127166+vitaliano-atlan@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:35:54 -0600 Subject: [PATCH 2/2] fix(lake-formation-tag-sync): route tables across candidate connections connection_map.json maps each database-name prefix to exactly one connection, so a feed whose tables live on several connections writes every row to one of them. Rows for tables on the other connections do not exist at the computed qualifiedName and the update-only import skips them. Let a connection_map.json value be a list of candidate database qualifiedNames as well as a single one. For a key with several candidates, look each table up across the candidates (active assets only) and write it to the one candidate that holds it. A table found under none or under several candidates is not written, and the run logs it with a count. A key with a single candidate keeps the existing path with no lookup, so existing maps are unchanged. An invalid value fails the run with an error. Also log connection keys that are missing from the map, per tag file, and update the SerdeTest fixture count for the added test records. Refs CSA-614 Co-Authored-By: Claude Opus 5.5 Signed-off-by: Vitaliano Vidaurre <122127166+vitaliano-atlan@users.noreply.github.com> --- .../kotlin/com/atlan/pkg/lftag/CSVProducer.kt | 198 ++++++++++++++---- .../atlan/pkg/lftag/LakeTagSynchronizer.kt | 42 +++- .../atlan/pkg/lftag/CSVProducerRoutingTest.kt | 106 ++++++++++ .../pkg/lftag/CSVProducerSchemaNameTest.kt | 2 +- .../com/atlan/pkg/lftag/CSVProducerTest.kt | 2 +- .../pkg/lftag/LakeTagSynchronizerTest.kt | 66 +++++- .../kotlin/com/atlan/pkg/lftag/SerdeTest.kt | 2 +- .../test/resources/lftag_association_1.json | 92 ++++++++ 8 files changed, 462 insertions(+), 48 deletions(-) create mode 100644 samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerRoutingTest.kt diff --git a/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/CSVProducer.kt b/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/CSVProducer.kt index 936a2acb68..cbfc99b2cb 100644 --- a/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/CSVProducer.kt +++ b/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/CSVProducer.kt @@ -7,15 +7,28 @@ import com.atlan.model.assets.Column import com.atlan.model.assets.Schema import com.atlan.model.assets.Table import com.atlan.pkg.Utils +import com.atlan.pkg.lftag.model.LFTableInfo import com.atlan.pkg.lftag.model.LFTagData import com.atlan.pkg.serde.csv.CSVWriter +import com.fasterxml.jackson.databind.JsonNode +/** + * Turns Lake Formation tag associations into an asset import CSV. + * + * @param connectionMap connection key (the database name prefix up to the first '_') to one or more + * candidate database qualifiedNames. A key with a single candidate is always written + * to that candidate. A key with several candidates writes each table to the one candidate + * that holds it, as reported by [tableLocator]. + * @param metadataMap LF tag key to custom metadata attribute + * @param tableLocator returns which of the given table qualifiedNames exist in Atlan; only called for keys + * with more than one candidate + */ class CSVProducer( - private val connectionMap: Map, + private val connectionMap: Map>, metadataMap: Map, + private val tableLocator: TableLocator? = null, ) { private val logger = Utils.getLogger(this.javaClass.name) - private val missingConnectionKeys = mutableSetOf() private val tagToMetadataMapper = TagToMetadataMapper(metadataMap) private val headerNames = setOf( @@ -24,6 +37,22 @@ class CSVProducer( Asset.QUALIFIED_NAME.atlanFieldName, ) + metadataMap.values.toSet() + /** Finds which table-level assets exist, by qualifiedName. */ + fun interface TableLocator { + /** + * @param qualifiedNames table-level qualifiedNames to look for + * @return the subset of [qualifiedNames] that exist as active assets + */ + fun findExisting(qualifiedNames: Collection): Set + } + + private data class TableEntry( + val info: LFTableInfo, + val connectionKey: String, + val schemaName: String, + val tableName: String, + ) + fun transform( tagData: LFTagData, fileName: String, @@ -33,6 +62,8 @@ class CSVProducer( CSVWriter(fileName).use { csv -> csv.writeHeader(headerNames) val start = System.currentTimeMillis() + val entries = mutableListOf() + val missingConnectionKeys = mutableSetOf() tagData.tableList.forEach { tableInfo -> val table = tableInfo.table val parts = table.databaseName.split('_', limit = 2) @@ -44,48 +75,139 @@ class CSVProducer( // By default the connection key prefix is stripped from the schema name; keep it when the // crawled schema name is the full database name. val schemaName = if (keepDatabasePrefix) table.databaseName else parts[1] - if (removeSchema) { - table.name = table.name.replace("$schemaName.", "") - } + val tableName = if (removeSchema) table.name.replace("$schemaName.", "") else table.name if (connectionKey in this.connectionMap) { - val schemaQualifiedName = "${connectionMap.getValue(connectionKey)}/$schemaName" - var row = - mutableMapOf( - Asset.QUALIFIED_NAME.atlanFieldName to schemaQualifiedName, - Asset.TYPE_NAME.atlanFieldName to Schema.TYPE_NAME, - Asset.NAME.atlanFieldName to schemaName, - ) - tagToMetadataMapper.getTagValues(tableInfo.lfTagOnDatabase, row) - csv.writeRecord(row) - val qualifiedName = "$schemaQualifiedName/${table.name}" - row = - mutableMapOf( - Asset.QUALIFIED_NAME.atlanFieldName to qualifiedName, - Asset.TYPE_NAME.atlanFieldName to Table.TYPE_NAME, - Asset.NAME.atlanFieldName to table.name, - ) - if (tableInfo.lfTagsOnTable != null) { - tagToMetadataMapper.getTagValues(tableInfo.lfTagsOnTable!!, row) - csv.writeRecord(row) - } - if (tableInfo.lfTagsOnColumn != null) { - tableInfo.lfTagsOnColumn?.forEach { column -> - val columnQualifiedName = "$qualifiedName/${column.name}" - row = - mutableMapOf( - Asset.QUALIFIED_NAME.atlanFieldName to columnQualifiedName, - Asset.TYPE_NAME.atlanFieldName to Column.TYPE_NAME, - Asset.NAME.atlanFieldName to column.name, - ) - tagToMetadataMapper.getTagValues(column.lfTags, row) - csv.writeRecord(row) - } - } + entries.add(TableEntry(tableInfo, connectionKey, schemaName, tableName)) } else { missingConnectionKeys.add(connectionKey) } } + val located = locateTables(entries) + var unresolved = 0 + var ambiguous = 0 + entries.forEach { entry -> + val candidates = connectionMap.getValue(entry.connectionKey) + val target = + if (candidates.size == 1) { + candidates[0] + } else { + val hits = candidates.filter { tableQualifiedName(it, entry) in located } + when (hits.size) { + 1 -> { + hits[0] + } + + 0 -> { + logger.warn { "Skipping ${entry.schemaName}/${entry.tableName}: not found under any connection mapped to '${entry.connectionKey}'." } + unresolved++ + null + } + + else -> { + logger.warn { "Skipping ${entry.schemaName}/${entry.tableName}: found under more than one connection mapped to '${entry.connectionKey}': $hits" } + ambiguous++ + null + } + } + } + if (target != null) { + writeRows(csv, target, entry) + } + } + if (missingConnectionKeys.isNotEmpty()) { + logger.warn { "Skipped every table whose connection key is not in the connection map: $missingConnectionKeys" } + } + if (unresolved > 0 || ambiguous > 0) { + logger.warn { "Tables not written because no mapped connection holds them: $unresolved; because several do: $ambiguous" } + } logger.info { "Total time taken: ${System.currentTimeMillis() - start} ms" } } } + + private fun tableQualifiedName( + databaseQualifiedName: String, + entry: TableEntry, + ) = "$databaseQualifiedName/${entry.schemaName}/${entry.tableName}" + + private fun locateTables(entries: List): Set { + val toLocate = + entries + .filter { connectionMap.getValue(it.connectionKey).size > 1 } + .flatMap { entry -> connectionMap.getValue(entry.connectionKey).map { tableQualifiedName(it, entry) } } + .toSet() + if (toLocate.isEmpty()) { + return emptySet() + } + val locator = + tableLocator + ?: throw IllegalStateException("A connection map key lists several connections, but no table locator was provided.") + return toLocate.chunked(LOOKUP_BATCH_SIZE).flatMap { locator.findExisting(it) }.toSet() + } + + private fun writeRows( + csv: CSVWriter, + databaseQualifiedName: String, + entry: TableEntry, + ) { + val tableInfo = entry.info + val schemaQualifiedName = "$databaseQualifiedName/${entry.schemaName}" + var row = + mutableMapOf( + Asset.QUALIFIED_NAME.atlanFieldName to schemaQualifiedName, + Asset.TYPE_NAME.atlanFieldName to Schema.TYPE_NAME, + Asset.NAME.atlanFieldName to entry.schemaName, + ) + tagToMetadataMapper.getTagValues(tableInfo.lfTagOnDatabase, row) + csv.writeRecord(row) + val qualifiedName = "$schemaQualifiedName/${entry.tableName}" + row = + mutableMapOf( + Asset.QUALIFIED_NAME.atlanFieldName to qualifiedName, + Asset.TYPE_NAME.atlanFieldName to Table.TYPE_NAME, + Asset.NAME.atlanFieldName to entry.tableName, + ) + if (tableInfo.lfTagsOnTable != null) { + tagToMetadataMapper.getTagValues(tableInfo.lfTagsOnTable!!, row) + csv.writeRecord(row) + } + tableInfo.lfTagsOnColumn?.forEach { column -> + val columnQualifiedName = "$qualifiedName/${column.name}" + row = + mutableMapOf( + Asset.QUALIFIED_NAME.atlanFieldName to columnQualifiedName, + Asset.TYPE_NAME.atlanFieldName to Column.TYPE_NAME, + Asset.NAME.atlanFieldName to column.name, + ) + tagToMetadataMapper.getTagValues(column.lfTags, row) + csv.writeRecord(row) + } + } + + companion object { + private const val LOOKUP_BATCH_SIZE = 500 + + /** + * Read connection_map.json, where each value is either one database qualifiedName or a list of them. + * + * @throws IllegalArgumentException if a value is neither a string nor a non-empty list of strings + */ + fun parseConnectionMap(root: JsonNode): Map> { + require(root.isObject) { "connection_map.json must be a JSON object." } + val result = mutableMapOf>() + root.fields().forEach { (key, value) -> + val candidates = + when { + value.isTextual -> listOf(value.asText()) + + value.isArray && value.size() > 0 && value.all { it.isTextual } -> value.map { it.asText() }.distinct() + + else -> throw IllegalArgumentException( + "connection_map.json key '$key' must map to a qualifiedName or a non-empty list of qualifiedNames.", + ) + } + result[key] = candidates + } + return result + } + } } diff --git a/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt b/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt index 96329cedcc..cf95363c9e 100644 --- a/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt +++ b/samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt @@ -5,6 +5,7 @@ package com.atlan.pkg.lftag import AssetImportCfg import LakeFormationTagSyncCfg import com.atlan.AtlanClient +import com.atlan.model.assets.Asset import com.atlan.pkg.PackageContext import com.atlan.pkg.Utils import com.atlan.pkg.aim.Importer @@ -13,6 +14,7 @@ import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper import com.fasterxml.jackson.module.kotlin.readValue import java.io.File +import java.util.stream.Collectors import kotlin.system.exitProcess /** @@ -57,15 +59,31 @@ object LakeTagSynchronizer { "", ) val tagFileNames = mutableListOf() - val connectionMap = mutableMapOf() + val connectionMap = mutableMapOf>() val metadataMap = mutableMapOf() files.forEach { fileName -> val nameWithExt = File(fileName).name when { - nameWithExt == CONNECTION_MAP_JSON -> connectionMap += mapper.readValue>(File(fileName)) - nameWithExt == METADATA_MAP_JSON -> metadataMap += mapper.readValue>(File(fileName)) - nameWithExt.startsWith(TAG_FILE_NAME_PREFIX) -> tagFileNames.add(fileName) - else -> logger.warn { "Skipping $fileName because it doesn't start with $TAG_FILE_NAME_PREFIX." } + nameWithExt == CONNECTION_MAP_JSON -> { + try { + connectionMap += CSVProducer.parseConnectionMap(mapper.readTree(File(fileName))) + } catch (e: IllegalArgumentException) { + logger.error { e.message } + return false + } + } + + nameWithExt == METADATA_MAP_JSON -> { + metadataMap += mapper.readValue>(File(fileName)) + } + + nameWithExt.startsWith(TAG_FILE_NAME_PREFIX) -> { + tagFileNames.add(fileName) + } + + else -> { + logger.warn { "Skipping $fileName because it doesn't start with $TAG_FILE_NAME_PREFIX." } + } } } if (connectionMap.isEmpty()) { @@ -80,7 +98,7 @@ object LakeTagSynchronizer { logger.error { "You must provide at least one json file containing Lake Tag data (name starting with $TAG_FILE_NAME_PREFIX)." } return false } - val csvProducer = CSVProducer(connectionMap, metadataMap) + val csvProducer = CSVProducer(connectionMap, metadataMap) { qualifiedNames -> findExisting(ctx.client, qualifiedNames) } tagFileNames.forEach { tagFileName -> logger.info { "Processing $tagFileName." } val csvFileName = "$outputDirectory${File.separator}${File(tagFileName).nameWithoutExtension}.csv" @@ -115,6 +133,18 @@ object LakeTagSynchronizer { return !anyFailure } + private fun findExisting( + client: AtlanClient, + qualifiedNames: Collection, + ): Set = + client.assets + .select() + .where(Asset.QUALIFIED_NAME.`in`(qualifiedNames)) + .pageSize(qualifiedNames.size.coerceAtLeast(1)) + .stream() + .map { it.qualifiedName } + .collect(Collectors.toSet()) + private fun createMissingEnums( client: AtlanClient, tagFileName: String, diff --git a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerRoutingTest.kt b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerRoutingTest.kt new file mode 100644 index 0000000000..534b0adfc1 --- /dev/null +++ b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerRoutingTest.kt @@ -0,0 +1,106 @@ +/* SPDX-License-Identifier: Apache-2.0 + Copyright 2026 Atlan Pte. Ltd. */ +package com.atlan.pkg.lftag + +import com.atlan.pkg.lftag.model.ColumnLFTag +import com.atlan.pkg.lftag.model.LFTable +import com.atlan.pkg.lftag.model.LFTableInfo +import com.atlan.pkg.lftag.model.LFTagData +import com.atlan.pkg.lftag.model.LFTagPair +import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper +import org.testng.annotations.Test +import java.io.File +import kotlin.io.path.createTempDirectory +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Covers routing each table to one of several candidate connections for a connection map key. + * Does not need a live tenant. + */ +class CSVProducerRoutingTest { + private val conn1 = "default/redshift/1111111111/db" + private val conn2 = "default/redshift/2222222222/db" + private val metadataMap = mapOf("security_classification" to "Data Privacy::Security Classification") + private val tag = listOf(LFTagPair("security_classification", listOf("public"), "123456789012")) + + private fun tagData(vararg tables: String) = + LFTagData( + tables.map { + LFTableInfo( + table = LFTable(databaseName = "dev_sales", name = it), + lfTagOnDatabase = tag, + lfTagsOnTable = tag, + lfTagsOnColumn = listOf(ColumnLFTag("col1", tag)), + ) + }, + ) + + private fun qualifiedNames( + connectionMap: Map>, + data: LFTagData, + locator: CSVProducer.TableLocator?, + ): List { + val file = File(createTempDirectory("lftag").toFile(), "out.csv") + CSVProducer(connectionMap, metadataMap, locator).transform(data, file.path) + val lines = file.readLines() + val qnIndex = lines.first().split(',').indexOfFirst { it.trim('"') == "qualifiedName" } + return lines.drop(1).map { it.split(',')[qnIndex].trim('"') } + } + + @Test + fun singleCandidateIsWrittenWithoutLookup() { + val locator = CSVProducer.TableLocator { throw AssertionError("lookup must not run for a single candidate") } + assertEquals( + listOf("$conn1/sales", "$conn1/sales/tbl1", "$conn1/sales/tbl1/col1"), + qualifiedNames(mapOf("dev" to listOf(conn1)), tagData("tbl1"), locator), + ) + } + + @Test + fun tableIsWrittenToTheOneCandidateThatHoldsIt() { + val existing = setOf("$conn1/sales/tbl1", "$conn2/sales/tbl2") + val locator = CSVProducer.TableLocator { qns -> qns.filter { it in existing }.toSet() } + assertEquals( + listOf( + "$conn1/sales", + "$conn1/sales/tbl1", + "$conn1/sales/tbl1/col1", + "$conn2/sales", + "$conn2/sales/tbl2", + "$conn2/sales/tbl2/col1", + ), + qualifiedNames(mapOf("dev" to listOf(conn1, conn2)), tagData("tbl1", "tbl2"), locator), + ) + } + + @Test + fun tableFoundNowhereIsNotWritten() { + val locator = CSVProducer.TableLocator { emptySet() } + assertTrue(qualifiedNames(mapOf("dev" to listOf(conn1, conn2)), tagData("tbl1"), locator).isEmpty()) + } + + @Test + fun tableFoundUnderSeveralCandidatesIsNotWritten() { + val locator = CSVProducer.TableLocator { qns -> qns.toSet() } + assertTrue(qualifiedNames(mapOf("dev" to listOf(conn1, conn2)), tagData("tbl1"), locator).isEmpty()) + } + + @Test + fun connectionMapAcceptsStringsAndLists() { + val root = jacksonObjectMapper().readTree("""{"dev": "$conn1", "val": ["$conn1", "$conn2", "$conn1"]}""") + assertEquals( + mapOf("dev" to listOf(conn1), "val" to listOf(conn1, conn2)), + CSVProducer.parseConnectionMap(root), + ) + } + + @Test + fun connectionMapRejectsOtherValues() { + val mapper = jacksonObjectMapper() + listOf("""{"dev": []}""", """{"dev": 1}""", """{"dev": ["$conn1", 2]}""", """["$conn1"]""").forEach { + assertFailsWith { CSVProducer.parseConnectionMap(mapper.readTree(it)) } + } + } +} diff --git a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerSchemaNameTest.kt b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerSchemaNameTest.kt index 6db86e76a1..605fb07599 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerSchemaNameTest.kt +++ b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerSchemaNameTest.kt @@ -19,7 +19,7 @@ import kotlin.test.assertTrue */ class CSVProducerSchemaNameTest { private val connectionQN = "default/athena/1234567890/example_db" - private val connectionMap = mapOf("prod" to connectionQN) + private val connectionMap = mapOf("prod" to listOf(connectionQN)) private val metadataMap = mapOf("security_classification" to "Data Privacy::Security Classification") private val tag = listOf(LFTagPair("security_classification", listOf("public"), "123456789012")) diff --git a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerTest.kt b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerTest.kt index bea17f2237..44505d91ba 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerTest.kt +++ b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/CSVProducerTest.kt @@ -17,7 +17,7 @@ class CSVProducerTest : PackageTest("csv") { @Test fun transform() { - val connectionMap = mapOf("dev" to "default/redshift/1687361820/wide_world_importers") + val connectionMap = mapOf("dev" to listOf("default/redshift/1687361820/wide_world_importers")) val metadataMap = mapOf( "business_domain" to "Data Domain::Business Domain", diff --git a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt index 817e41bd13..b43c3f34ff 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt +++ b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizerTest.kt @@ -26,6 +26,7 @@ import java.io.File import java.nio.file.Paths import kotlin.io.path.copyTo import kotlin.test.assertEquals +import kotlin.test.assertNull private const val PUBLIC = "public" @@ -36,12 +37,19 @@ private const val FULL_HISTORY = "fullhistory" class LakeTagSynchronizerTest : PackageTest("lts") { override val logger = Utils.getLogger(this.javaClass.name) private val c1 = makeUnique("c1") + private val c2 = makeUnique("c2") private var connectionQualifiedName = "" + private var connection2QualifiedName = "" private var schemaGuid = "" private var tableGuid = "" private var columnGuid = "" private var viewGuid = "" private var viewColumnGuid = "" + private var schema2Guid = "" + private var table2Guid = "" + private var column2Guid = "" + private var dup1Guid = "" + private var dup2Guid = "" private val cm1 = makeUnique("cm") private val enum1 = makeUnique("enum") @@ -60,6 +68,11 @@ class LakeTagSynchronizerTest : PackageTest("lts") { .build() .save(client) .block() + Connection + .creator(client, c2, connectorType) + .build() + .save(client) + .block() } private fun createAssets() { @@ -85,10 +98,43 @@ class LakeTagSynchronizerTest : PackageTest("lts") { viewGuid = response.getResult(vw).guid viewColumnGuid = response.getResult(vwColumn).guid } + // A second connection under the same connection map key: sch2.tbl2 exists only here, + // and sch.dup1 exists on both connections, so it cannot be routed. + val connection2 = Connection.findByName(client, c2, connectorType)[0]!! + connection2QualifiedName = connection2.qualifiedName + AssetBatch(client, 20).use { batch -> + val db = Database.creator(databaseName, connection2.qualifiedName).build() + batch.add(db) + val sch2 = Schema.creator("sch2", db).build() + batch.add(sch2) + val tbl2 = Table.creator("tbl2", sch2).build() + batch.add(tbl2) + val column2 = Column.creator("col2", tbl2, 1).build() + batch.add(column2) + val sch = Schema.creator("sch", db).build() + batch.add(sch) + val dup2 = Table.creator("dup1", sch).build() + batch.add(dup2) + val dup1 = Table.creator("dup1", "$connectionQualifiedName/$databaseName/sch").build() + batch.add(dup1) + val response = batch.flush() + schema2Guid = response.getResult(sch2).guid + table2Guid = response.getResult(tbl2).guid + column2Guid = response.getResult(column2).guid + dup1Guid = response.getResult(dup1).guid + dup2Guid = response.getResult(dup2).guid + } } private fun createConnectionMap() { - val connectionMap = mapOf("dev" to "$connectionQualifiedName/$databaseName") + val connectionMap = + mapOf( + "dev" to + listOf( + "$connectionQualifiedName/$databaseName", + "$connection2QualifiedName/$databaseName", + ), + ) Paths .get(testDirectory, directoryPrefix, CONNECTION_MAP_JSON) .toFile() @@ -156,6 +202,7 @@ class LakeTagSynchronizerTest : PackageTest("lts") { override fun teardown() { removeConnection(c1, connectorType) + removeConnection(c2, connectorType) client.typeDefs.purge(client.customMetadataCache.getSidForName(cm1)) EnumDef.purge(client, enum1) } @@ -203,6 +250,23 @@ class LakeTagSynchronizerTest : PackageTest("lts") { assertEquals(PUBLIC, attribute1) } + @Test + fun validateTableRoutedToSecondConnection() { + val schema = Schema.get(client, schema2Guid) + assertEquals(PUBLIC, schema.getCustomMetadata(cm1, attr1)) + val table = Table.get(client, table2Guid) + assertEquals(PUBLIC, table.getCustomMetadata(cm1, attr1)) + assertEquals(FULL_HISTORY, table.getCustomMetadata(cm1, attr3)) + val column = Column.get(client, column2Guid) + assertEquals(PUBLIC, column.getCustomMetadata(cm1, attr1)) + } + + @Test + fun validateAmbiguousTableNotTagged() { + assertNull(Table.get(client, dup1Guid).getCustomMetadata(cm1, attr1)) + assertNull(Table.get(client, dup2Guid).getCustomMetadata(cm1, attr1)) + } + @Test fun validateFilesCreated() { validateFilesExist( diff --git a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/SerdeTest.kt b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/SerdeTest.kt index 2bf14f4963..c0c71791c2 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/SerdeTest.kt +++ b/samples/packages/lake-formation-tag-sync/src/test/kotlin/com/atlan/pkg/lftag/SerdeTest.kt @@ -172,6 +172,6 @@ class SerdeTest { fun whenDeserializableTagDataThenSuccess() { val jsonString: String = File("./src/test/resources/lftag_association_1.json").readText(Charsets.UTF_8) val tagData = mapper.readValue(jsonString, LFTagData::class.java) - assertEquals(3, tagData.tableList.size) + assertEquals(6, tagData.tableList.size) } } diff --git a/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json b/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json index e1820b98c2..9bdfcf8989 100644 --- a/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json +++ b/samples/packages/lake-formation-tag-sync/src/test/resources/lftag_association_1.json @@ -128,6 +128,98 @@ ] } ] + }, + { + "Table": { + "CatalogId": "614518280298", + "DatabaseName": "dev_sch2", + "Name": "sch2.tbl2" + }, + "LFTagOnDatabase": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + } + ], + "LFTagsOnTable": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + }, + { + "CatalogId": "614518280298", + "TagKey": "data_load_method", + "TagValues": [ + "fullhistory" + ] + } + ], + "LFTagsOnColumns": [ + { + "Name": "col2", + "LFTags": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + } + ] + } + ] + }, + { + "Table": { + "CatalogId": "614518280298", + "DatabaseName": "dev_sch", + "Name": "sch.dup1" + }, + "LFTagOnDatabase": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + } + ], + "LFTagsOnTable": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + }, + { + "CatalogId": "614518280298", + "TagKey": "data_load_method", + "TagValues": [ + "fullhistory" + ] + } + ], + "LFTagsOnColumns": [ + { + "Name": "dcol1", + "LFTags": [ + { + "CatalogId": "614518280298", + "TagKey": "security_classification", + "TagValues": [ + "public" + ] + } + ] + } + ] } ] }