APP-4791: Added support for APITokenConnectionAdmin workflow package

atlanhq · Jan 28, 2025 · 0620cd5 · 0620cd5
1 parent 63f5314
commit 0620cd5
Show file tree

Hide file tree

Showing 5 changed files with 168 additions and 0 deletions.
diff --git a/pyatlan/model/enums.py b/pyatlan/model/enums.py
@@ -2227,11 +2227,13 @@ class WorkflowPackage(str, Enum):
     TERADATA_MINER = "atlan-teradata-miner"
     THOUGHTSPOT = "atlan-thoughtspot"
     TRINO = "atlan-trino"
+    # CSA packages
     ASSET_IMPORT = "csa-asset-import"
     ASSET_EXPORT_BASIC = "csa-asset-export-basic"
     RELATIONAL_ASSETS_BUILDER = "csa-relational-assets-builder"
     LINEAGE_BUILDER = "csa-lineage-builder"
     LINEAGE_GENERATOR = "csa-lineage-generator"
+    API_TOKEN_CONNECTION_ADMIN = "csa-api-token-connection-admin"  # noqa: S105
 
 
 class AssetInputHandling(str, Enum):

diff --git a/pyatlan/model/packages/__init__.py b/pyatlan/model/packages/__init__.py
@@ -1,4 +1,5 @@
 # flake8: noqa
+from .api_token_connection_admin import APITokenConnectionAdmin
 from .asset_export_basic import AssetExportBasic
 from .asset_import import AssetImport
 from .big_query_crawler import BigQueryCrawler
@@ -45,4 +46,5 @@
     "OracleCrawler",
     "LineageBuilder",
     "LineageGenerator",
+    "APITokenConnectionAdmin",
 ]
diff --git a/pyatlan/model/packages/api_token_connection_admin.py b/pyatlan/model/packages/api_token_connection_admin.py
@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from pyatlan.model.enums import WorkflowPackage
+from pyatlan.model.packages.base.custom_package import AbstractCustomPackage
+from pyatlan.model.workflow import WorkflowMetadata
+
+
+class APITokenConnectionAdmin(AbstractCustomPackage):
+    """
+    Base configuration for a new API token connection admin package.
+    """
+
+    _NAME = "api-token-connection-admin"
+    _PACKAGE_NAME = f"@csa/{_NAME}"
+    _PACKAGE_PREFIX = WorkflowPackage.API_TOKEN_CONNECTION_ADMIN.value
+    _PACKAGE_ICON = "http://assets.atlan.com/assets/ph-key-light.svg"
+    _PACKAGE_LOGO = "http://assets.atlan.com/assets/ph-key-light.svg"
+
+    def config(
+        self, connection_qualified_name: str, api_token_guid: str
+    ) -> APITokenConnectionAdmin:
+        """
+        Set up the API token connection admin with the specified configuration.
+
+        :param connection_qualified_name: connection qualified name
+        to which you want to add the API token as a connection admin.
+        :param api_token_guid: guid of the API token
+
+        :returns: package, with the specified configuration.
+        """
+        self._parameters.append(
+            {"name": "connection_qualified_name", "value": connection_qualified_name}
+        )
+        self._parameters.append({"name": "api_token_guid", "value": api_token_guid})
+        return self
+
+    def _get_metadata(self) -> WorkflowMetadata:
+        return WorkflowMetadata(
+            labels={
+                "orchestration.atlan.com/certified": "true",
+                "orchestration.atlan.com/source": "atlan",
+                "orchestration.atlan.com/sourceCategory": "utility",
+                "orchestration.atlan.com/type": "custom",
+                "orchestration.atlan.com/preview": "true",
+                "orchestration.atlan.com/verified": "true",
+                "package.argoproj.io/installer": "argopm",
+                "package.argoproj.io/name": f"a-t-rcsas-l-a-s-h{self._NAME}",
+                "package.argoproj.io/registry": "httpsc-o-l-o-ns-l-a-s-hs-l-a-s-hpackages.atlan.com",
+                "orchestration.atlan.com/atlan-ui": "true",
+            },
+            annotations={
+                "orchestration.atlan.com/allowSchedule": "true",
+                "orchestration.atlan.com/categories": "kotlin,utility",
+                "orchestration.atlan.com/dependentPackage": "",
+                "orchestration.atlan.com/docsUrl": f"https://solutions.atlan.com/{self._NAME}/",
+                "orchestration.atlan.com/emoji": "\U0001f680",
+                "orchestration.atlan.com/icon": self._PACKAGE_ICON,
+                "orchestration.atlan.com/logo": self._PACKAGE_LOGO,
+                "orchestration.atlan.com/name": "API Token Connection Admin",
+                "package.argoproj.io/author": "Atlan CSA",
+                "package.argoproj.io/description": "Assigns an API token as a connection admin for an existing connection.",  # noqa
+                "package.argoproj.io/homepage": f"https://packages.atlan.com/-/web/detail/{self._PACKAGE_NAME}",
+                "package.argoproj.io/keywords": "[\"kotlin\",\"utility\"]",  # fmt: skip
+                "package.argoproj.io/name": self._PACKAGE_NAME,
+                "package.argoproj.io/parent": ".",
+                "package.argoproj.io/registry": "https://packages.atlan.com",
+                "package.argoproj.io/repository": "git+https://github.com/atlanhq/marketplace-packages.git",
+                "package.argoproj.io/support": "[email protected]",
+                "orchestration.atlan.com/atlanName": f"csa-{self._NAME}-{self._epoch}",
+            },
+            name=f"csa-{self._NAME}-{self._epoch}",
+            namespace="default",
+        )
diff --git a/tests/unit/data/package_requests/api_token_connection_admin.json b/tests/unit/data/package_requests/api_token_connection_admin.json
@@ -0,0 +1,76 @@
+{
+    "metadata": {
+      "annotations": {
+        "orchestration.atlan.com/allowSchedule": "true",
+        "orchestration.atlan.com/categories": "kotlin,utility",
+        "orchestration.atlan.com/dependentPackage": "",
+        "orchestration.atlan.com/docsUrl": "https://solutions.atlan.com/api-token-connection-admin/",
+        "orchestration.atlan.com/emoji": "🚀",
+        "orchestration.atlan.com/icon": "http://assets.atlan.com/assets/ph-key-light.svg",
+        "orchestration.atlan.com/logo": "http://assets.atlan.com/assets/ph-key-light.svg",
+        "orchestration.atlan.com/name": "API Token Connection Admin",
+        "package.argoproj.io/author": "Atlan CSA",
+        "package.argoproj.io/description": "Assigns an API token as a connection admin for an existing connection.",
+        "package.argoproj.io/homepage": "https://packages.atlan.com/-/web/detail/@csa/api-token-connection-admin",
+        "package.argoproj.io/keywords": "[\"kotlin\",\"utility\"]",
+        "package.argoproj.io/name": "@csa/api-token-connection-admin",
+        "package.argoproj.io/parent": ".",
+        "package.argoproj.io/registry": "https://packages.atlan.com",
+        "package.argoproj.io/repository": "git+https://github.com/atlanhq/marketplace-packages.git",
+        "package.argoproj.io/support": "[email protected]",
+        "orchestration.atlan.com/atlanName": "csa-api-token-connection-admin-123456"
+      },
+      "labels": {
+        "orchestration.atlan.com/certified": "true",
+        "orchestration.atlan.com/source": "atlan",
+        "orchestration.atlan.com/sourceCategory": "utility",
+        "orchestration.atlan.com/type": "custom",
+        "orchestration.atlan.com/preview": "true",
+        "orchestration.atlan.com/verified": "true",
+        "package.argoproj.io/installer": "argopm",
+        "package.argoproj.io/name": "a-t-rcsas-l-a-s-hapi-token-connection-admin",
+        "package.argoproj.io/registry": "httpsc-o-l-o-ns-l-a-s-hs-l-a-s-hpackages.atlan.com",
+        "orchestration.atlan.com/atlan-ui": "true"
+      },
+      "name": "csa-api-token-connection-admin-123456",
+      "namespace": "default"
+    },
+    "spec": {
+      "entrypoint": "main",
+      "templates": [
+        {
+          "name": "main",
+          "dag": {
+            "tasks": [
+              {
+                "name": "run",
+                "arguments": {
+                  "parameters": [
+                    {
+                      "name": "connection_qualified_name",
+                      "value": "default/snowflake/1234567890"
+                    },
+                    {
+                      "name": "api_token_guid",
+                      "value": "92588c67-5ddf-4a45-8b5c-dd92f4b84e99"
+                    }
+                  ]
+                },
+                "templateRef": {
+                  "name": "csa-api-token-connection-admin",
+                  "template": "main",
+                  "clusterScope": true
+                }
+              }
+            ]
+          }
+        }
+      ],
+      "workflowMetadata": {
+        "annotations": {
+          "package.argoproj.io/name": "@csa/api-token-connection-admin"
+        }
+      }
+    },
+    "payload": []
+}
diff --git a/tests/unit/test_packages.py b/tests/unit/test_packages.py
@@ -8,6 +8,7 @@
 from pyatlan.model.assets.core import Asset
 from pyatlan.model.enums import AssetDeltaHandling, AssetInputHandling, AssetRemovalType
 from pyatlan.model.packages import (
+    APITokenConnectionAdmin,
     AssetExportBasic,
     AssetImport,
     BigQueryCrawler,
@@ -98,6 +99,7 @@
 LINEAGE_BUILDER_ADLS = "lineage_builder_adls.json"
 LINEAGE_GENERATOR_DEFAULT = "lineage_generator_default.json"
 LINEAGE_GENERATOR_FULL = "lineage_generator_full.json"
+API_TOKEN_CONNECTION_ADMIN = "api_token_connection_admin.json"
 
 
 class NonSerializable:
@@ -1554,6 +1556,19 @@ def test_lineage_generator_nt(mock_package_env):
     assert request_json == load_json(LINEAGE_GENERATOR_FULL)
 
 
+def test_api_token_connection_admin(mock_package_env):
+    token_connection_admin = (
+        APITokenConnectionAdmin()
+        .config(
+            connection_qualified_name="default/snowflake/1234567890",
+            api_token_guid="92588c67-5ddf-4a45-8b5c-dd92f4b84e99",
+        )
+        .to_workflow()
+    )
+    request_json = loads(token_connection_admin.json(by_alias=True, exclude_none=True))
+    assert request_json == load_json(API_TOKEN_CONNECTION_ADMIN)
+
+
 @pytest.mark.parametrize(
     "test_assets",
     [