Merge pull request #625 from atlanhq/APP-6565

APP-6565: Added support for dynamic extension of `AtlanConnectorType` for custom connectors

Author: Aryamanz29

pyatlan/cache/connection_cache.py

@@ -175,8 +175,18 @@ def __init__(
         elif isinstance(connection, str):
             tokens = connection.split("/")
             if len(tokens) > 1:
-                self.type = AtlanConnectorType(tokens[0]).value  # type: ignore[call-arg]
-                self.name = connection[len(tokens[0]) + 1 :]  # noqa
+                # Try enum conversion; fallback to custom connector if it fails
+                try:
+                    self.type = AtlanConnectorType(tokens[0]).value  # type: ignore[call-arg]
+                    self.name = connection[len(tokens[0]) + 1 :]  # noqa
+                except ValueError:
+                    custom_connector = AtlanConnectorType.CREATE_CUSTOM(
+                        # Ensure the enum name is converted to UPPER_SNAKE_CASE from kebab-case
+                        name=tokens[0].replace("-", "_").upper(),
+                        value=tokens[0],
+                    )
+                    self.type = custom_connector.value
+                    self.name = connection[len(tokens[0]) + 1 :]
 
     def __hash__(self):
         return hash((self.name, self.type))

pyatlan/model/enums.py

@@ -140,9 +140,21 @@ class AtlanConnectionCategory(str, Enum):
     CUSTOM = "custom"
 
 
-class AtlanConnectorType(str, Enum):
+class AtlanConnectorType(str, Enum, metaclass=utils.ExtendableEnumMeta):
     category: AtlanConnectionCategory
 
+    @classmethod
+    def get_values(cls):
+        return [member.value for member in cls._member_map_.values()]
+
+    @classmethod
+    def get_names(cls):
+        return list(cls._member_map_.keys())
+
+    @classmethod
+    def get_items(cls):
+        return [(name, member.value) for name, member in cls._member_map_.items()]
+
     @classmethod
     def _get_connector_type_from_qualified_name(
         cls, qualified_name: str
@@ -152,12 +164,17 @@ def _get_connector_type_from_qualified_name(
             raise ValueError(
                 f"Qualified name '{qualified_name}' does not contain enough segments."
             )
-        connector_type_key = tokens[1].upper()
-        # Check if the connector_type_key exists in AtlanConnectorType
+
+        connector_value = tokens[1]
+        # Ensure the enum name is converted to UPPER_SNAKE_CASE from kebab-case
+        connector_type_key = tokens[1].replace("-", "_").upper()
+
+        # Check if the connector_type_key exists in AtlanConnectorType;
+        # if so, return it directly. Otherwise, it may be a custom type.
         if connector_type_key not in AtlanConnectorType.__members__:
-            raise ValueError(
-                f"Could not determine AtlanConnectorType from '{qualified_name}'; "
-                f"'{connector_type_key}' is not a valid connector type."
+            return AtlanConnectorType.CREATE_CUSTOM(
+                name=connector_type_key,
+                value=connector_value,
             )
         return AtlanConnectorType[connector_type_key]
 
@@ -169,6 +186,12 @@ def __new__(
         obj.category = category
         return obj
 
+    @classmethod
+    def CREATE_CUSTOM(
+        cls, name: str, value: str, category=AtlanConnectionCategory.CUSTOM
+    ) -> "AtlanConnectorType":
+        return cls.add_value(name, value, category)
+
     def to_qualified_name(self):
         return f"default/{self.value}/{int(utils.get_epoch_timestamp())}"
 
@@ -193,14 +216,22 @@ def get_connector_name(
         fields = qualified_name.split("/")
         if len(fields) != qualified_name_len:
             raise ValueError(err)
+
+        connector_value = fields[1]
+        # Try enum conversion; fallback to custom connector if it fails
         try:
-            connector_name = AtlanConnectorType(fields[1]).value  # type:ignore
-            if attribute_name != "connection_qualified_name":
-                connection_qn = f"{fields[0]}/{fields[1]}/{fields[2]}"
-                return connection_qn, connector_name
-            return connector_name
-        except ValueError as e:
-            raise ValueError(err) from e
+            connector_name = AtlanConnectorType(connector_value).value  # type: ignore
+        except ValueError:
+            custom_connection = AtlanConnectorType.CREATE_CUSTOM(
+                # Ensure the enum name is converted to UPPER_SNAKE_CASE from kebab-case
+                name=connector_value.replace("-", "_").upper(),
+                value=connector_value,
+            )
+            connector_name = custom_connection.value
+        if attribute_name != "connection_qualified_name":
+            connection_qn = f"{fields[0]}/{fields[1]}/{fields[2]}"
+            return connection_qn, connector_name
+        return connector_name
 
     SNOWFLAKE = ("snowflake", AtlanConnectionCategory.WAREHOUSE)
     TABLEAU = ("tableau", AtlanConnectionCategory.BI)

pyatlan/utils.py

@@ -10,7 +10,7 @@
 import time
 from contextvars import ContextVar
 from datetime import datetime
-from enum import Enum
+from enum import Enum, EnumMeta
 from functools import reduce, wraps
 from typing import Any, Dict, List, Mapping, Optional
 
@@ -471,3 +471,21 @@ def validate_single_required_field(field_names: List[str], values: List[Any]):
         raise ValueError(
             f"Only one of the following parameters are allowed: {', '.join(names)}"
         )
+
+
+class ExtendableEnumMeta(EnumMeta):
+    def __init__(cls, name, bases, namespace):
+        super().__init__(name, bases, namespace)
+        cls._additional_members = {}
+
+    def add_value(cls, name, value, category=None):
+        member = str.__new__(cls, value)
+        member._name_ = name
+        member._value_ = value
+        member.category = category
+
+        cls._member_map_[name] = member
+        cls._value2member_map_[value] = member
+        cls._member_names_.append(name)
+        cls._additional_members[name] = member
+        return member

tests/integration/connection_test.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
+from typing import Generator
+
 import pytest
 
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.assets import Connection
-from pyatlan.model.enums import AtlanConnectorType
-from tests.integration.client import TestId
+from pyatlan.model.enums import AtlanConnectionCategory, AtlanConnectorType
+from tests.integration.client import TestId, delete_asset
 
 MODULE_NAME = TestId.make_unique("CONN")
 
@@ -24,6 +26,29 @@ def create_connection(
     )
 
 
+@pytest.fixture(scope="module")
+def custom_connection(client: AtlanClient) -> Generator[Connection, None, None]:
+    CUSTOM_CONNECTOR_TYPE = AtlanConnectorType.CREATE_CUSTOM(
+        name=f"{MODULE_NAME}_NAME",
+        value=f"{MODULE_NAME}_type",
+        category=AtlanConnectionCategory.API,
+    )
+    result = create_connection(
+        client=client, name=MODULE_NAME, connector_type=CUSTOM_CONNECTOR_TYPE
+    )
+    yield result
+    # TODO: proper connection delete workflow
+    delete_asset(client, guid=result.guid, asset_type=Connection)
+
+
+def test_custom_connection(custom_connection: Connection):
+    assert custom_connection.name == MODULE_NAME
+    assert custom_connection.connector_name == f"{MODULE_NAME}_type"
+    assert custom_connection.qualified_name
+    assert f"default/{MODULE_NAME}_type" in custom_connection.qualified_name
+    assert AtlanConnectorType[f"{MODULE_NAME}_NAME"].value == f"{MODULE_NAME}_type"
+
+
 def test_invalid_connection(client: AtlanClient):
     with pytest.raises(
         ValueError, match="One of admin_user, admin_groups or admin_roles is required"

tests/unit/model/a_d_l_s_object_test.py

@@ -134,7 +134,7 @@ def test_create():
         ),
         (
             ADLS_OBJECT_NAME,
-            "default/adls-invalid/production",
+            "default/adls/invalid/production",
             "abc",
             ADLS_CONTAINER_NAME,
             ADLS_CONTAINER_QUALIFIED_NAME,

tests/unit/model/connection_test.py

@@ -6,7 +6,7 @@
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.client.token import TokenClient
 from pyatlan.model.assets import Connection
-from pyatlan.model.enums import AtlanConnectorType
+from pyatlan.model.enums import AtlanConnectionCategory, AtlanConnectorType
 from tests.unit.model.constants import CONNECTION_NAME, CONNECTION_QUALIFIED_NAME
 
 
@@ -174,6 +174,69 @@ def test_create(
     assert sut.admin_roles == set(admin_roles)
 
 
+@pytest.mark.parametrize(
+    "name, connector_type, admin_users, admin_groups, admin_roles",
+    [
+        (
+            CONNECTION_NAME,
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="FOO", value="foo", category=AtlanConnectionCategory.BI
+            ),
+            ["ernest"],
+            [],
+            [],
+        ),
+        (
+            CONNECTION_NAME,
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="BAR", value="bar", category=AtlanConnectionCategory.API
+            ),
+            [],
+            ["ernest"],
+            [],
+        ),
+        (
+            CONNECTION_NAME,
+            AtlanConnectorType.CREATE_CUSTOM(
+                name="BAZ", value="baz", category=AtlanConnectionCategory.WAREHOUSE
+            ),
+            [],
+            [],
+            ["ernest"],
+        ),
+    ],
+)
+def test_creator_with_custom_type(
+    name: str,
+    connector_type: AtlanConnectorType,
+    admin_users: List[str],
+    admin_groups: List[str],
+    admin_roles: List[str],
+    mock_role_cache,
+    mock_user_cache,
+    mock_group_cache,
+):
+    mock_role_cache.validate_idstrs
+    mock_user_cache.validate_names
+    mock_group_cache.validate_aliases
+
+    sut = Connection.create(
+        name=name,
+        connector_type=connector_type,
+        admin_users=admin_users,
+        admin_groups=admin_groups,
+        admin_roles=admin_roles,
+    )
+
+    assert sut.name == name
+    assert sut.qualified_name
+    assert sut.qualified_name[:20] == connector_type.to_qualified_name()[:20]
+    assert sut.connector_name == connector_type.value
+    assert sut.admin_users == set(admin_users)
+    assert sut.admin_groups == set(admin_groups)
+    assert sut.admin_roles == set(admin_roles)
+
+
 @pytest.mark.parametrize(
     "qualified_name, name, message",
     [