Merge pull request #219 from atlanhq/DVX-61

Add code to support custom package toolkit

HISTORY.md

@@ -1,3 +1,17 @@
+## 1.9.0 (January 16, 2024)
+
+### New Features
+
+- Add ability to update certificate, announcement for GlossaryTerm and GlossaryCategory
+- Add `create` method for `ColumnProcess`
+- Always include sort by `GUID` as final criteria in `IndexSearch`
+- (Experimental) Add classes to support custom package generation
+
+### QOL Improvements
+
+- Add an additional parameter to `create` method of `ADLSObject`
+- Add type checking to AtlanClient save and other methods
+
 ## 1.8.4 (January 4, 2024)
 
 ### New Features

pyatlan/__init__.py

@@ -2,13 +2,11 @@
 # Copyright 2022 Atlan Pte. Ltd.
 import logging.config
 import os
-from logging import NullHandler
 
 from pyatlan.utils import REQUEST_ID_FILTER
 
-LOGGER = logging.getLogger(__name__)
-LOGGER.addHandler(NullHandler())
 if os.path.exists("logging.conf"):
     logging.config.fileConfig("logging.conf")
+LOGGER = logging.getLogger(__name__)
 for handler in LOGGER.handlers:
     handler.addFilter(REQUEST_ID_FILTER)

pyatlan/client/atlan.py

@@ -34,6 +34,7 @@
 from pyatlan.client.constants import PARSE_QUERY, UPLOAD_IMAGE
 from pyatlan.client.credential import CredentialClient
 from pyatlan.client.group import GroupClient
+from pyatlan.client.impersonate import ImpersonationClient
 from pyatlan.client.role import RoleClient
 from pyatlan.client.search_log import SearchLogClient
 from pyatlan.client.token import TokenClient
@@ -64,7 +65,13 @@
 from pyatlan.model.typedef import TypeDef, TypeDefResponse
 from pyatlan.model.user import AtlanUser, UserMinimalResponse, UserResponse
 from pyatlan.multipart_data_generator import MultipartDataGenerator
-from pyatlan.utils import API, AuthorizationFilter, HTTPStatus, RequestIdAdapter
+from pyatlan.utils import (
+    API,
+    APPLICATION_ENCODED_FORM,
+    AuthorizationFilter,
+    HTTPStatus,
+    RequestIdAdapter,
+)
 
 request_id_var = ContextVar("request_id", default=None)
 
@@ -133,6 +140,7 @@ class AtlanClient(BaseSettings):
     _typedef_client: Optional[TypeDefClient] = PrivateAttr(default=None)
     _token_client: Optional[TokenClient] = PrivateAttr(default=None)
     _user_client: Optional[UserClient] = PrivateAttr(default=None)
+    _impersonate_client: Optional[ImpersonationClient] = PrivateAttr(default=None)
 
     class Config:
         env_prefix = "atlan_"
@@ -218,6 +226,12 @@ def asset(self) -> AssetClient:
             self._asset_client = AssetClient(client=self)
         return self._asset_client
 
+    @property
+    def impersonate(self) -> ImpersonationClient:
+        if self._impersonate_client is None:
+            self._impersonate_client = ImpersonationClient(client=self)
+        return self._impersonate_client
+
     @property
     def token(self) -> TokenClient:
         if self._token_client is None:
@@ -236,6 +250,9 @@ def user(self) -> UserClient:
             self._user_client = UserClient(client=self)
         return self._user_client
 
+    def update_headers(self, header: dict[str, str]):
+        self._session.headers.update(header)
+
     def _call_api_internal(self, api, path, params, binary_data=None):
         token = request_id_var.set(str(uuid.uuid4()))
         try:
@@ -351,6 +368,8 @@ def _create_params(
                 params["data"] = request_obj.json(
                     by_alias=True, exclude_unset=exclude_unset
                 )
+            elif api.consumes == APPLICATION_ENCODED_FORM:
+                params["data"] = request_obj
             else:
                 params["data"] = json.dumps(request_obj)
         return params

pyatlan/client/constants.py

@@ -3,6 +3,7 @@
 # Based on original code from https://github.com/apache/atlas (under Apache-2.0 license)
 from pyatlan.utils import (
     API,
+    APPLICATION_ENCODED_FORM,
     APPLICATION_JSON,
     APPLICATION_OCTET_STREAM,
     MULTIPART_FORM_DATA,
@@ -110,6 +111,15 @@
     TOKENS_API, HTTPMethod.DELETE, HTTPStatus.OK, endpoint=EndPoint.HERACLES
 )
 
+GET_TOKEN = API(
+    "/auth/realms/default/protocol/openid-connect/token",
+    HTTPMethod.POST,
+    HTTPStatus.OK,
+    endpoint=EndPoint.IMPERSONATION,
+    consumes=APPLICATION_ENCODED_FORM,
+    produces=APPLICATION_ENCODED_FORM,
+)
+
 ENTITY_API = "entity/"
 PREFIX_ATTR = "attr:"
 PREFIX_ATTR_ = "attr_"

pyatlan/client/impersonate.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 Atlan Pte. Ltd.
+
+import logging
+import os
+from typing import NamedTuple
+
+from pyatlan.client.common import ApiCaller
+from pyatlan.client.constants import GET_TOKEN
+from pyatlan.errors import AtlanError, ErrorCode
+from pyatlan.model.response import AccessTokenResponse
+
+LOGGER = logging.getLogger(__name__)
+
+
+class ClientInfo(NamedTuple):
+    client_id: str
+    client_secret: str
+
+
+class ImpersonationClient:
+    """
+    This class can be used for impersonating users as part of Atlan automations (if desired).
+    Note: this will only work when run as part of Atlan's packaged workflow ecosystem (running in the cluster back-end).
+    """
+
+    def __init__(self, client: ApiCaller):
+        if not isinstance(client, ApiCaller):
+            raise ErrorCode.INVALID_PARAMETER_TYPE.exception_with_parameters(
+                "client", "ApiCaller"
+            )
+        self._client = client
+
+    def user(self, user_id: str) -> str:
+        """
+        Retrieves a bearer token that impersonates the provided user.
+
+        :param user_id: unique identifier of the user to impersonate
+        :returns: a bearer token that impersonates the provided user
+        :raises AtlanError: on any API communication issue
+        """
+        client_info = self._get_client_info()
+        credentials = {
+            "grant_type": "client_credentials",
+            "client_id": client_info.client_id,
+            "client_secret": client_info.client_secret,
+        }
+
+        LOGGER.debug("Getting token with client id and secret")
+        try:
+            raw_json = self._client._call_api(GET_TOKEN, request_obj=credentials)
+            argo_token = AccessTokenResponse(**raw_json).access_token
+        except AtlanError as atlan_err:
+            raise ErrorCode.UNABLE_TO_ESCALATE.exception_with_parameters() from atlan_err
+        LOGGER.debug("Getting token with subject token")
+        try:
+            user_credentials = {
+                "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+                "client_id": client_info.client_id,
+                "client_secret": client_info.client_secret,
+                "subject_token": argo_token,
+                "requested_subject": user_id,
+            }
+            raw_json = self._client._call_api(GET_TOKEN, request_obj=user_credentials)
+            return AccessTokenResponse(**raw_json).access_token
+        except AtlanError as atlan_err:
+            raise ErrorCode.UNABLE_TO_IMPERSONATE.exception_with_parameters() from atlan_err
+
+    def _get_client_info(self) -> ClientInfo:
+        client_id = os.getenv("CLIENT_ID")
+        client_secret = os.getenv("CLIENT_SECRET")
+        if not client_id or not client_secret:
+            raise ErrorCode.MISSING_CREDENTIALS.exception_with_parameters()
+        client_info = ClientInfo(client_id=client_id, client_secret=client_secret)
+        return client_info
+
+    def escalate(self) -> str:
+        """
+        Escalate to a privileged user on a short-term basis.
+        Note: this is only possible from within the Atlan tenant, and only when given the appropriate credentials.
+
+        :returns: a short-lived bearer token with escalated privileges
+        :raises AtlanError: on any API communication issue
+        """
+        client_info = self._get_client_info()
+        credentials = {
+            "grant_type": "client_credentials",
+            "client_id": client_info.client_id,
+            "client_secret": client_info.client_secret,
+            "scope": "openid",
+        }
+        try:
+            raw_json = self._client._call_api(GET_TOKEN, request_obj=credentials)
+            return AccessTokenResponse(**raw_json).access_token
+        except AtlanError as atlan_err:
+            raise ErrorCode.UNABLE_TO_ESCALATE.exception_with_parameters() from atlan_err

pyatlan/client/token.py

@@ -98,6 +98,21 @@ def get_by_id(self, client_id: str) -> Optional[ApiToken]:
                 return response.records[0]
         return None
 
+    @validate_arguments
+    def get_by_guid(self, guid: str) -> Optional[ApiToken]:
+        """
+        Retrieves the API token with a unique ID (GUID) that exactly matches the provided string.
+
+        :param guid: unique identifier by which to retrieve the API token
+        :returns: the API token whose clientId matches the provided string, or None if there is none
+        """
+        if response := self.get(
+            offset=0, limit=5, post_filter='{"id":"' + guid + '"}', sort="createdAt"
+        ):
+            if response.records and len(response.records) >= 1:
+                return response.records[0]
+        return None
+
     @validate_arguments
     def create(
         self,

pyatlan/errors.py

@@ -484,6 +484,13 @@ class ErrorCode(Enum):
         "Please double-check your method keyword arguments.",
         InvalidRequestError,
     )
+    MISSING_CREDENTIALS = (
+        400,
+        "ATLAN-PYTHON-400-056",
+        "Missing privileged credentials to impersonate users.",
+        "You must have both CLIENT_ID and CLIENT_SECRET configured to be able to impersonate users.",
+        InvalidRequestError,
+    )
     AUTHENTICATION_PASSTHROUGH = (
         401,
         "ATLAN-PYTHON-401-000",
@@ -535,6 +542,20 @@ class ErrorCode(Enum):
         "Check the details of the server's message to correct your request.",
         PermissionError,
     )
+    UNABLE_TO_ESCALATE = (
+        403,
+        "ATLAN-PYTHON-403-001",
+        "Unable to escalate to a privileged user.",
+        "Check the details of your configured privileged credentials.",
+        PermissionError,
+    )
+    UNABLE_TO_IMPERSONATE = (
+        403,
+        "ATLAN-PYTHON-403-002",
+        "Unable to impersonate requested user.",
+        "Check the details of your configured privileged credentials and the user you requested to impersonate.",
+        PermissionError,
+    )
     NOT_FOUND_PASSTHROUGH = (
         404,
         "ATLAN-PYTHON-404-000",

pyatlan/logging.conf

@@ -13,13 +13,13 @@ handlers=consoleHandler
 
 [logger_pyatlan]
 level=DEBUG
-handlers=consoleHandler,fileHandler,jsonHandler
+handlers=fileHandler,jsonHandler
 qualname=pyatlan
 propagate=0
 
 [logger_urllib3]
 level=DEBUG
-handlers=consoleHandler,fileHandler,jsonHandler
+handlers=fileHandler,jsonHandler
 qualname=urllib3
 propagate=0
 

pyatlan/model/response.py

@@ -85,3 +85,14 @@ def assets_partially_updated(self, asset_type: Type[A]) -> list[A]:
                 if isinstance(asset, asset_type)
             ]
         return []
+
+
+class AccessTokenResponse(AtlanObject):
+    access_token: str
+    expires_in: int
+    refresh_expires_in: int
+    refresh_token: str
+    token_type: str
+    not_before_policy: Optional[int] = Field(default=None)
+    session_state: str
+    scope: str