Skip to content

fix(lake-formation-tag-sync): match views when syncing LF tags #2804

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
vitaliano-atlan wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

fix(lake-formation-tag-sync): match views when syncing LF tags #2804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fa6f939
fix(lake-formation-tag-sync): match views when syncing LF tags
vitaliano-atlan Oct 6, 2026
413eb37
fix(lake-formation-tag-sync): route tables across candidate connections
vitaliano-atlan Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
198 changes: 160 additions & 38 deletions samples/packages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/CSVProducer.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String>,
private val connectionMap: Map<String, List<String>>,
metadataMap: Map<String, String>,
private val tableLocator: TableLocator? = null,
) {
private val logger = Utils.getLogger(this.javaClass.name)
private val missingConnectionKeys = mutableSetOf<String>()
private val tagToMetadataMapper = TagToMetadataMapper(metadataMap)
private val headerNames =
setOf(
Expand All @@ -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<String>): Set<String>
}

private data class TableEntry(
val info: LFTableInfo,
val connectionKey: String,
val schemaName: String,
val tableName: String,
)

fun transform(
tagData: LFTagData,
fileName: String,
Expand All @@ -33,6 +62,8 @@ class CSVProducer(
CSVWriter(fileName).use { csv ->
csv.writeHeader(headerNames)
val start = System.currentTimeMillis()
val entries = mutableListOf<TableEntry>()
val missingConnectionKeys = mutableSetOf<String>()
tagData.tableList.forEach { tableInfo ->
val table = tableInfo.table
val parts = table.databaseName.split('_', limit = 2)
Expand All @@ -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<TableEntry>): Set<String> {
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<String, List<String>> {
require(root.isObject) { "connection_map.json must be a JSON object." }
val result = mutableMapOf<String, List<String>>()
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
}
}
}
45 changes: 39 additions & 6 deletions ...ckages/lake-formation-tag-sync/src/main/kotlin/com/atlan/pkg/lftag/LakeTagSynchronizer.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

/**
Expand Down Expand Up @@ -57,15 +59,31 @@ object LakeTagSynchronizer {
"",
)
val tagFileNames = mutableListOf<String>()
val connectionMap = mutableMapOf<String, String>()
val connectionMap = mutableMapOf<String, List<String>>()
val metadataMap = mutableMapOf<String, String>()
files.forEach { fileName ->
val nameWithExt = File(fileName).name
when {
nameWithExt == CONNECTION_MAP_JSON -> connectionMap += mapper.readValue<Map<String, String>>(File(fileName))
nameWithExt == METADATA_MAP_JSON -> metadataMap += mapper.readValue<Map<String, String>>(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<Map<String, String>>(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()) {
Expand All @@ -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"
Expand All @@ -90,6 +108,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(
Expand All @@ -112,6 +133,18 @@ object LakeTagSynchronizer {
return !anyFailure
}

private fun findExisting(
client: AtlanClient,
qualifiedNames: Collection<String>,
): Set<String> =
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,
Expand Down
Loading
Loading