From ca8946921aa61989ce47995f816f0773d7e337e6 Mon Sep 17 00:00:00 2001 From: Christopher Grote Date: Tue, 26 Sep 2023 17:33:14 +0100 Subject: [PATCH 1/2] Bump to latest snapshot Signed-off-by: Christopher Grote --- gradle.properties | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gradle.properties b/gradle.properties index d381307..4b983f6 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1,2 +1,2 @@ kotlin.code.style=official -version=0.2.1-SNAPSHOT +version=0.2.2-SNAPSHOT From 253a0f4fd1e3776b1297699d2edf91faf2f0c7af Mon Sep 17 00:00:00 2001 From: Christopher Grote Date: Tue, 26 Sep 2023 19:01:53 +0100 Subject: [PATCH 2/2] Adds custom package to set API token as a connection admin Signed-off-by: Christopher Grote --- api-token-connection-admin/build.gradle.kts | 15 +++ .../main/kotlin/ApiTokenConnectionAdmin.kt | 97 +++++++++++++++++++ common/src/main/kotlin/Utils.kt | 84 +++++++++------- gradle.properties | 2 +- settings.gradle.kts | 1 + 5 files changed, 165 insertions(+), 34 deletions(-) create mode 100644 api-token-connection-admin/build.gradle.kts create mode 100644 api-token-connection-admin/src/main/kotlin/ApiTokenConnectionAdmin.kt diff --git a/api-token-connection-admin/build.gradle.kts b/api-token-connection-admin/build.gradle.kts new file mode 100644 index 0000000..0895206 --- /dev/null +++ b/api-token-connection-admin/build.gradle.kts @@ -0,0 +1,15 @@ +val jarPath = "$rootDir/jars" + +plugins { + id("atlan-kotlin-sample") +} + +dependencies { + implementation(project(":common")) +} + +tasks { + jar { + destinationDirectory.set(file(jarPath)) + } +} diff --git a/api-token-connection-admin/src/main/kotlin/ApiTokenConnectionAdmin.kt b/api-token-connection-admin/src/main/kotlin/ApiTokenConnectionAdmin.kt new file mode 100644 index 0000000..c990680 --- /dev/null +++ b/api-token-connection-admin/src/main/kotlin/ApiTokenConnectionAdmin.kt @@ -0,0 +1,97 @@ +/* SPDX-License-Identifier: Apache-2.0 */ +/* Copyright 2023 Atlan Pte. Ltd. */ +import com.atlan.Atlan +import com.atlan.exception.AtlanException +import com.atlan.model.assets.Asset +import com.atlan.model.assets.Connection +import com.atlan.model.core.AssetMutationResponse +import mu.KotlinLogging +import kotlin.system.exitProcess + +private val log = KotlinLogging.logger {} + +/** + * Actually run the logic to add the API token as a connection admin. + */ +fun main() { + Utils.setClient() + Utils.setWorkflowOpts() + + val connectionQN = Utils.reuseConnection("CONNECTION_QUALIFIED_NAME") + val apiTokenName = Utils.getEnvVar("API_TOKEN_NAME", "") + + if (connectionQN == "" || apiTokenName == "") { + log.error("Missing required parameter - you must provide BOTH a connection and the name of an API token.") + exitProcess(4) + } + + val apiTokenId = getIdForToken(apiTokenName) + val connection = getConnectionWithAdmins(connectionQN) + addTokenAsConnectionAdmin(connection, apiTokenId) +} + +/** + * Retrieve the API token's pseudo-username, that can be used anywhere a username can be used. + * + * @param apiTokenName name of the API token for which to fetch the pseudo-username + * @return the pseudo-username of the API token + */ +fun getIdForToken(apiTokenName: String): String { + log.info("Looking up API token: {}", apiTokenName) + val token = Atlan.getDefaultClient().apiTokens.get(apiTokenName) + if (token == null) { + log.error("Unable to find any API token with the name: {}", apiTokenName) + exitProcess(5) + } + return "service-account-${token.clientId}" +} + +/** + * Retrieve the connection with its existing admins. + * + * @param connectionQN qualifiedName of the connection + * @return the connection with its existing admins + */ +fun getConnectionWithAdmins(connectionQN: String): Asset { + log.info("Looking up connection details: {}", connectionQN) + val found = Connection.select() + .where(Connection.QUALIFIED_NAME.eq(connectionQN)) + .includeOnResults(Connection.ADMIN_USERS) + .stream() + .findFirst() + if (found.isEmpty) { + log.error("Unable to find the specified connection: {}", connectionQN) + exitProcess(6) + } + return found.get() +} + +/** + * Actually add the token as a connection admin, appending it to any pre-existing + * connection admins (rather than replacing). + * + * @param connection the connection to add the API token to, with its existing admin users present + * @param apiToken the API token to append as a connection admin + */ +fun addTokenAsConnectionAdmin(connection: Asset, apiToken: String) { + log.info("Adding API token {} as connection admin for: {}", apiToken, connection.qualifiedName) + val existingAdmins = connection.adminUsers + try { + val response = connection.trimToRequired() + .adminUsers(existingAdmins) + .adminUser(apiToken) + .build() + .save() + when (val result = response?.getMutation(connection)) { + AssetMutationResponse.MutationType.UPDATED -> log.info(" ... successfully updated the connection with API token as a new admin.") + AssetMutationResponse.MutationType.NOOP -> log.info(" ... API token is already an admin on the connection - no changes made.") + AssetMutationResponse.MutationType.CREATED -> log.error(" ... somehow created the connection - that should not have happened.") + AssetMutationResponse.MutationType.DELETED -> log.error(" ... somehow deleted the connection - that should not have happened.") + else -> { + log.warn("Unexpected connection change result: {}", result) + } + } + } catch (e: AtlanException) { + log.error("Unable to add the API token as a connection admin.", e) + } +} diff --git a/common/src/main/kotlin/Utils.kt b/common/src/main/kotlin/Utils.kt index 4c35551..dc3c358 100644 --- a/common/src/main/kotlin/Utils.kt +++ b/common/src/main/kotlin/Utils.kt @@ -131,42 +131,60 @@ object Utils { */ fun createOrReuseConnection(varForAction: String, varForReuse: String, varForCreate: String): String { val action = getEnvVar(varForAction, "REUSE") - val connectionQN: String - if (action == "REUSE") { - val providedConnectionQN = getEnvVar(varForReuse, "") + return if (action == "REUSE") { + reuseConnection(varForReuse) + } else { + createConnection(varForCreate) + } + } + + /** + * Create a connection using the details provided through the provided environment variable. + * + * @param varWithConnectionString name of the environment variable containing a full connection object, as a string + * @return the qualifiedName of the connection that is created, or an empty string if no connection details exist in the environment variable + */ + fun createConnection(varWithConnectionString: String): String { + val connectionString = getEnvVar(varWithConnectionString, "") + return if (connectionString != "") { + log.info("Attempting to create new connection...") try { - log.info("Attempting to reuse connection: {}", providedConnectionQN) - Connection.get(Atlan.getDefaultClient(), providedConnectionQN, false) - } catch (e: NotFoundException) { - log.error("Unable to find connection with the provided qualifiedName: {}", providedConnectionQN, e) - exitProcess(1) + val toCreate = Atlan.getDefaultClient().readValue(connectionString, Connection::class.java) + .toBuilder() + .guid("-${ThreadLocalRandom.current().nextLong(0, Long.MAX_VALUE - 1)}") + .build() + val response = toCreate.save().block() + response.getResult(toCreate).qualifiedName + } catch (e: IOException) { + log.error("Unable to deserialize the connection details: {}", connectionString, e) + exitProcess(2) + } catch (e: IllegalArgumentException) { + log.error("Unable to deserialize the connection details: {}", connectionString, e) + exitProcess(2) + } catch (e: AtlanException) { + log.error("Unable to create connection: {}", connectionString, e) + exitProcess(3) } - connectionQN = providedConnectionQN } else { - val connectionString = getEnvVar(varForCreate, "") - connectionQN = if (connectionString != "") { - log.info("Attempting to create new connection...") - try { - val toCreate = Atlan.getDefaultClient().readValue(connectionString, Connection::class.java) - .toBuilder() - .guid("-${ThreadLocalRandom.current().nextLong(0, Long.MAX_VALUE - 1)}") - .build() - val response = toCreate.save().block() - response.getResult(toCreate).qualifiedName - } catch (e: IOException) { - log.error("Unable to deserialize the connection details: {}", connectionString, e) - exitProcess(2) - } catch (e: IllegalArgumentException) { - log.error("Unable to deserialize the connection details: {}", connectionString, e) - exitProcess(2) - } catch (e: AtlanException) { - log.error("Unable to create connection: {}", connectionString, e) - exitProcess(3) - } - } else { - "" - } + "" + } + } + + /** + * Validate the provided connection exists, and if so return its qualifiedName. + * + * @param varWithConnectionQN name of the environment variable containing a connection's qualifiedName + * @return the qualifiedName of the connection, so long as it exists, otherwise an empty string + */ + fun reuseConnection(varWithConnectionQN: String): String { + val providedConnectionQN = getEnvVar(varWithConnectionQN, "") + return try { + log.info("Attempting to reuse connection: {}", providedConnectionQN) + Connection.get(Atlan.getDefaultClient(), providedConnectionQN, false) + providedConnectionQN + } catch (e: NotFoundException) { + log.error("Unable to find connection with the provided qualifiedName: {}", providedConnectionQN, e) + "" } - return connectionQN } } diff --git a/gradle.properties b/gradle.properties index 4b983f6..03aeef8 100644 --- a/gradle.properties +++ b/gradle.properties @@ -1,2 +1,2 @@ kotlin.code.style=official -version=0.2.2-SNAPSHOT +version=0.3.0-SNAPSHOT diff --git a/settings.gradle.kts b/settings.gradle.kts index f4e818b..b76a5b3 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -11,6 +11,7 @@ plugins { rootProject.name = "atlan-kotlin-samples" include("common") +include("api-token-connection-admin") include("duplicate-detector") include("migration-assistant") include("openapi-spec-loader") \ No newline at end of file