atlan-application-sdk 0.1.1rc26__py3-none-any.whl → 0.1.1rc28__py3-none-any.whl

application_sdk/clients/atlan_client.py

@@ -0,0 +1,70 @@
+from typing import Optional
+
+from pyatlan.client.atlan import AtlanClient
+
+from application_sdk.common.error_codes import ClientError
+from application_sdk.constants import (
+    ATLAN_API_KEY,
+    ATLAN_API_TOKEN_GUID,
+    ATLAN_BASE_URL,
+    ATLAN_CLIENT_ID,
+    ATLAN_CLIENT_SECRET,
+)
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_client(
+    base_url: Optional[str] = None,
+    api_key: Optional[str] = None,
+    api_token_guid: Optional[str] = None,
+) -> AtlanClient:
+    """
+    Returns an authenticated AtlanClient instance using provided parameters or environment variables.
+
+    Selects authentication method based on the presence of parameters or environment variables and validates the required configuration.
+    In general, the use of environment variables is recommended. Any parameters specified will override the environment variables.
+
+    Args:
+    base_url: Atlan base URL (overrides ATLAN_BASE_URL)
+    api_key: Atlan API key (overrides ATLAN_API_KEY)
+    api_token_guid: API token GUID (overrides API_TOKEN_GUID)
+    """
+    # Resolve final values (parameters override env vars)
+    final_token_guid = api_token_guid or ATLAN_API_TOKEN_GUID
+    final_base_url = base_url or ATLAN_BASE_URL
+    final_api_key = api_key or ATLAN_API_KEY
+
+    # Priority 1: Token-based auth (recommended for production)
+    if final_token_guid:
+        if final_base_url or final_api_key:  # Check original params, not env vars
+            logger.warning(
+                "Token auth takes precedence - ignoring base_url/api_key parameters as well as ATLAN_BASE_URL and ATLAN_API_KEY environment variables."
+            )
+        return _get_client_from_token(final_token_guid)
+
+    # Priority 2: API key + base URL auth
+    if not final_base_url:
+        raise ClientError(
+            "ATLAN_BASE_URL is required (via parameter or environment variable)"
+        )
+    if not final_api_key:
+        raise ClientError(
+            "ATLAN_API_KEY is required (via parameter or environment variable)"
+        )
+
+    logger.info("Using API key-based authentication")
+    return AtlanClient(base_url=final_base_url, api_key=final_api_key)
+
+
+def _get_client_from_token(api_token_guid: str):
+    if not ATLAN_CLIENT_ID:
+        raise ClientError(
+            f"{ClientError.AUTH_CONFIG_ERROR}: Environment variable CLIENT_ID is required when API_TOKEN_GUID is set."
+        )
+    if not ATLAN_CLIENT_SECRET:
+        raise ClientError(
+            f"{ClientError.AUTH_CONFIG_ERROR}: Environment variable CLIENT_SECRET is required when API_TOKEN_GUID is set."
+        )
+    return AtlanClient.from_token_guid(guid=api_token_guid)

application_sdk/constants.py

@@ -27,11 +27,14 @@ from dotenv import load_dotenv
 
 load_dotenv(dotenv_path=".env")
 
+# Static Constants
+LOCAL_ENVIRONMENT = "local"
+
 # Application Constants
 #: Name of the application, used for identification
 APPLICATION_NAME = os.getenv("ATLAN_APPLICATION_NAME", "default")
 #: Name of the deployment, used to distinguish between different deployments of the same application
-DEPLOYMENT_NAME = os.getenv("ATLAN_DEPLOYMENT_NAME", "local")
+DEPLOYMENT_NAME = os.getenv("ATLAN_DEPLOYMENT_NAME", LOCAL_ENVIRONMENT)
 #: Host address for the application's HTTP server
 APP_HOST = str(os.getenv("ATLAN_APP_HTTP_HOST", "localhost"))
 #: Port number for the application's HTTP server
@@ -46,8 +49,6 @@ APP_DASHBOARD_PORT = int(os.getenv("ATLAN_APP_DASHBOARD_PORT", "8000"))
 SQL_SERVER_MIN_VERSION = os.getenv("ATLAN_SQL_SERVER_MIN_VERSION")
 #: Path to the SQL queries directory
 SQL_QUERIES_PATH = os.getenv("ATLAN_SQL_QUERIES_PATH", "app/sql")
-#: Whether to use local development mode (used for instance to fetch secrets from the local state store)
-LOCAL_DEVELOPMENT = os.getenv("ATLAN_LOCAL_DEVELOPMENT", "false").lower() == "true"
 
 # Output Path Constants
 #: Output path format for workflows (example: objectstore://bucket/artifacts/apps/{application_name}/workflows/{workflow_id}/{workflow_run_id})
@@ -216,3 +217,10 @@ TRACES_FILE_NAME = "traces.parquet"
 ENABLE_OBSERVABILITY_DAPR_SINK = (
     os.getenv("ATLAN_ENABLE_OBSERVABILITY_DAPR_SINK", "false").lower() == "true"
 )
+
+# atlan_client configuration (non ATLAN_ prefix are rooted in pyatlan SDK, to be revisited)
+ATLAN_API_TOKEN_GUID = os.getenv("API_TOKEN_GUID")
+ATLAN_BASE_URL = os.getenv("ATLAN_BASE_URL")
+ATLAN_API_KEY = os.getenv("ATLAN_API_KEY")
+ATLAN_CLIENT_ID = os.getenv("CLIENT_ID")
+ATLAN_CLIENT_SECRET = os.getenv("CLIENT_SECRET")

application_sdk/inputs/objectstore.py

@@ -2,7 +2,7 @@
 
 import json
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Union
 
 import orjson
 from dapr.clients import DaprClient
@@ -27,7 +27,7 @@ class ObjectStoreInput:
         cls,
         operation: str,
         metadata: Dict[str, str],
-        data: Optional[bytes] = None,
+        data: Union[bytes, str] = "",
         object_store_name: str = DEPLOYMENT_OBJECT_STORE_NAME,
     ) -> bytes:
         """
@@ -172,7 +172,7 @@ class ObjectStoreInput:
         try:
             # this takes care of listing from all type of storage - local as well as object stores
             metadata = {"prefix": prefix, "fileName": prefix} if prefix else {}
-            data = json.dumps({"prefix": prefix}).encode("utf-8") if prefix else None
+            data = json.dumps({"prefix": prefix}).encode("utf-8") if prefix else ""
 
             response_data = cls._invoke_dapr_binding(
                 operation=cls.OBJECT_LIST_OPERATION,
@@ -233,7 +233,7 @@ class ObjectStoreInput:
         """
         try:
             metadata = {"key": file_path, "fileName": file_path}
-            data = json.dumps({"key": file_path}).encode("utf-8") if file_path else None
+            data = json.dumps({"key": file_path}).encode("utf-8") if file_path else ""
 
             response_data = cls._invoke_dapr_binding(
                 operation=cls.OBJECT_GET_OPERATION,

application_sdk/inputs/parquet.py

@@ -1,4 +1,5 @@
 import glob
+import os
 from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
 
 from application_sdk.inputs import Input
@@ -41,27 +42,32 @@ class ParquetInput(Input):
         self.input_prefix = input_prefix
         self.file_names = file_names
 
-    async def download_files(self, remote_file_path: str) -> Optional[str]:
+    async def download_files(self, local_file_path: str) -> Optional[str]:
         """Read a file from the object store.
 
         Args:
-            remote_file_path (str): Path to the remote file in object store.
+            local_file_path (str): Path to the local file in the temp directory.
 
         Returns:
             Optional[str]: Path to the downloaded local file.
         """
-        parquet_files = glob.glob(remote_file_path)
+        parquet_files = glob.glob(local_file_path)
         if not parquet_files:
             if self.input_prefix:
                 logger.info(
-                    f"Reading file from object store: {remote_file_path} from {self.input_prefix}"
-                )
-                ObjectStoreInput.download_files_from_object_store(
-                    self.input_prefix, remote_file_path
+                    f"Reading file from object store: {local_file_path} from {self.input_prefix}"
                 )
+                if os.path.isdir(local_file_path):
+                    ObjectStoreInput.download_files_from_object_store(
+                        self.input_prefix, local_file_path
+                    )
+                else:
+                    ObjectStoreInput.download_file_from_object_store(
+                        self.input_prefix, local_file_path
+                    )
             else:
                 raise ValueError(
-                    f"No parquet files found in {remote_file_path} and no input prefix provided"
+                    f"No parquet files found in {local_file_path} and no input prefix provided"
                 )
 
     async def get_dataframe(self) -> "pd.DataFrame":

application_sdk/inputs/secretstore.py

@@ -8,9 +8,10 @@ from typing import Any, Dict
 from dapr.clients import DaprClient
 
 from application_sdk.constants import (
+    DEPLOYMENT_NAME,
     DEPLOYMENT_SECRET_PATH,
     DEPLOYMENT_SECRET_STORE_NAME,
-    LOCAL_DEVELOPMENT,
+    LOCAL_ENVIRONMENT,
     SECRET_STORE_NAME,
 )
 from application_sdk.observability.logger_adaptor import get_logger
@@ -41,7 +42,7 @@ class SecretStoreInput:
         Returns:
             Dict with processed secret data
         """
-        if LOCAL_DEVELOPMENT:
+        if DEPLOYMENT_NAME == LOCAL_ENVIRONMENT:
             return {}
 
         try:

application_sdk/outputs/atlan_storage.py

@@ -82,7 +82,6 @@ class AtlanStorageOutput:
 
             logger.debug(f"Successfully migrated: {file_path}")
             return file_path, True, ""
-
         except Exception as e:
             error_msg = str(e)
             logger.error(f"Failed to migrate file {file_path}: {error_msg}")

application_sdk/outputs/objectstore.py

@@ -1,6 +1,7 @@
 """Object store interface for the application."""
 
 import os
+import shutil
 
 from dapr.clients import DaprClient
 from temporalio import activity
@@ -15,6 +16,22 @@ activity.logger = logger
 class ObjectStoreOutput:
     OBJECT_CREATE_OPERATION = "create"
 
+    @staticmethod
+    def _cleanup_local_path(path: str) -> None:
+        """Remove a file or directory (recursively). Ignores if doesn't exist.
+        Args:
+            path (str): The path to the file or directory to remove.
+        """
+        try:
+            if os.path.isfile(path) or os.path.islink(path):
+                os.remove(path)
+            elif os.path.isdir(path):
+                shutil.rmtree(path)
+        except FileNotFoundError:
+            pass  # ignore if the file or directory doesn't exist
+        except Exception as e:
+            logger.warning(f"Error cleaning up {path}: {str(e)}")
+
     @classmethod
     async def push_file_to_object_store(
         cls,
@@ -62,6 +79,7 @@ class ObjectStoreOutput:
                     f"Error pushing file {relative_path} to object store: {str(e)}"
                 )
                 raise e
+            cls._cleanup_local_path(file_path)
 
     @classmethod
     async def push_files_to_object_store(

application_sdk/outputs/secretstore.py

@@ -3,7 +3,7 @@
 import uuid
 from typing import Any, Dict
 
-from application_sdk.constants import LOCAL_DEVELOPMENT
+from application_sdk.constants import DEPLOYMENT_NAME, LOCAL_ENVIRONMENT
 from application_sdk.inputs.statestore import StateType
 from application_sdk.outputs.statestore import StateStoreOutput
 
@@ -26,7 +26,7 @@ class SecretStoreOutput:
             >>> SecretStoreOutput.save_secret({"username": "admin", "password": "password"})
             "1234567890"
         """
-        if LOCAL_DEVELOPMENT:
+        if DEPLOYMENT_NAME == LOCAL_ENVIRONMENT:
             # NOTE: (development) temporary solution to store the credentials in the state store.
             # In production, dapr doesn't support creating secrets.
             credential_guid = str(uuid.uuid4())

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "0.1.1rc26"
+__version__ = "0.1.1rc28"

{atlan_application_sdk-0.1.1rc26.dist-info → atlan_application_sdk-0.1.1rc28.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc26
+Version: 0.1.1rc28
 Summary: Atlan Application SDK is a Python library for developing applications on the Atlan Platform
 Project-URL: Repository, https://github.com/atlanhq/application-sdk
 Project-URL: Documentation, https://github.com/atlanhq/application-sdk/README.md
@@ -26,12 +26,12 @@ Requires-Dist: fastapi[standard]>=0.115.0
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: opentelemetry-exporter-otlp>=1.27.0
 Requires-Dist: psutil>=7.0.0
-Requires-Dist: pyatlan>=7.1.1
+Requires-Dist: pyatlan>=7.1.6
 Requires-Dist: pydantic>=2.10.6
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: uvloop>=0.21.0; sys_platform != 'win32'
 Provides-Extra: daft
-Requires-Dist: getdaft[sql]>=0.4.12; extra == 'daft'
+Requires-Dist: daft[sql]>=0.4.12; extra == 'daft'
 Provides-Extra: iam-auth
 Requires-Dist: boto3>=1.38.6; extra == 'iam-auth'
 Provides-Extra: iceberg

{atlan_application_sdk-0.1.1rc26.dist-info → atlan_application_sdk-0.1.1rc28.dist-info}/RECORD

@@ -1,6 +1,6 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
-application_sdk/constants.py,sha256=txVsXgN9aBsDXw0Dw9XqYbdTsLTVP5qTS26I1IKpVCA,9111
-application_sdk/version.py,sha256=fqQ1fw6udcOHy2GZF5O3JVl7ADaHX-iU2cezF5M4P04,88
+application_sdk/constants.py,sha256=Npmw6ojj731faIE3BMeXE5bucxWFkAD13cYx5cQ9PZw,9302
+application_sdk/version.py,sha256=dAt_fXyCvt-5sLv5kCnHWbaTLVTYgdCQ1V7EoOwVJDU,88
 application_sdk/worker.py,sha256=2fLjuKNJafhaQXrHzmxXYO22F4ZSc0igMjoxXVNBFfk,6167
 application_sdk/activities/__init__.py,sha256=EH5VTHcfGykIX7V1HsG0J1Z-1FbJEXTQOET0HdzFDjU,9519
 application_sdk/activities/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -15,6 +15,7 @@ application_sdk/application/__init__.py,sha256=wCaAsoQ-wIYtfeG6hNwwLi8PcufjwMBjm
 application_sdk/application/metadata_extraction/sql.py,sha256=ohpV4qZ92uKRlH7I_8G67ocnWkZJAZCU_7XdvqYPiN4,7966
 application_sdk/clients/__init__.py,sha256=C9T84J7V6ZumcoWJPAxdd3tqSmbyciaGBJn-CaCCny0,1341
 application_sdk/clients/atlan_auth.py,sha256=MQznmvVKrlOT_Tp232W4UrOupRrx9Dx9zQm3n1R7kD8,8938
+application_sdk/clients/atlan_client.py,sha256=f2-Uk5KiPIDJEhGkfYctA_f3CwoVB_mWNBMVvxeLuY4,2684
 application_sdk/clients/sql.py,sha256=tW89SHuuWdU5jv8lDUP5AUCEpR2CF_5TyUvYDCBHses,17880
 application_sdk/clients/temporal.py,sha256=kiB9W2e3x6gGtFC9e7vqTIBNGbt7nyS6p7j2hLKdsgI,22364
 application_sdk/clients/utils.py,sha256=zLFOJbTr_6TOqnjfVFGY85OtIXZ4FQy_rquzjaydkbY,779
@@ -45,9 +46,9 @@ application_sdk/handlers/sql.py,sha256=qyLbTG9HZYX5_PAsvoSKANVcK53JS3sLUrqXMgXr1
 application_sdk/inputs/__init__.py,sha256=_d-cUhcDyoJTJR3PdQkC831go6VDw9AM6Bg7-qm3NHI,1900
 application_sdk/inputs/iceberg.py,sha256=xiv1kNtVx1k0h3ZJbJeXjZwdfBGSy9j9orYP_AyCYlI,2756
 application_sdk/inputs/json.py,sha256=J1CVz0YGQHDyq840TyoBHE7Baua2yIFHzsrybiZbeWk,6079
-application_sdk/inputs/objectstore.py,sha256=z7BaEpVzoiscZ2Tdq9vPBWj4MZb3LjygE2dd1jn1ZRo,8788
-application_sdk/inputs/parquet.py,sha256=rrDI_OAktZIpWMIJIvyNoqrDg-lH2t_4NnB_x6JmX30,6101
-application_sdk/inputs/secretstore.py,sha256=3JieHdRXEOATWwJW4jkuRvcd7EDUHR0akSXhqtSM0Q0,3906
+application_sdk/inputs/objectstore.py,sha256=mOC5Bv10tZg9y4oikff_QViHSWZ_FHiotI9xdrlAIok,8781
+application_sdk/inputs/parquet.py,sha256=Hsi6Nz_KwxFMB6DcHSQov5y_hRkoeN7e4xfpYwogveo,6346
+application_sdk/inputs/secretstore.py,sha256=WPGvsPMLUaARvXA6JSa4uHp7tQcCW4NMIIUDQyM3F-I,3946
 application_sdk/inputs/sql_query.py,sha256=1EREgea6kKNaMIyX2HLJgbJ07rtAgLasd9NyvDcdZok,10636
 application_sdk/inputs/statestore.py,sha256=Ib87toiYzSEg7_7C0o6CeS-c0vsjaJqiIEIKwiwXAcE,2954
 application_sdk/observability/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -58,13 +59,13 @@ application_sdk/observability/traces_adaptor.py,sha256=0eQJPN-tYA_dV8D3uEa5ZiX9g
 application_sdk/observability/utils.py,sha256=MKEpT0WYtpATUgLgJDkGQaAP_t-jpDYMUKDfEvr8Phg,2448
 application_sdk/observability/decorators/observability_decorator.py,sha256=JNrWNXT5W4klmlAc5b8C3_VBjDu0PI64W2ptr7LMzk4,8110
 application_sdk/outputs/__init__.py,sha256=HX8VcN22xyrkoRWdjQj8TrC5dEUG7cPzOcvJhlprqAs,8415
-application_sdk/outputs/atlan_storage.py,sha256=V4fXslizhpsOg7_on_a_RM8IEAl6wa3fqd8BcceXjsg,6030
+application_sdk/outputs/atlan_storage.py,sha256=HQLbuyOZQC-GxYAiCVJakIJizTWy926tdMGOHvaBlD8,6029
 application_sdk/outputs/eventstore.py,sha256=fEW1oo9ncKylF5uLXTgjyRQZCrj72Gta3CaeduSX8nw,5548
 application_sdk/outputs/iceberg.py,sha256=IGtj5WDgqLu6vzDEvw5DLsKsjm29Krto3AHvWpemr0A,5311
 application_sdk/outputs/json.py,sha256=xF-8mY3BZRRejip4s9npIUuFaAxgFmBQVaLMkrI_iCI,14117
-application_sdk/outputs/objectstore.py,sha256=CfEI13tgp-WyQAEDyBelUZULUzAxHMxrdZifn66mIKo,3771
+application_sdk/outputs/objectstore.py,sha256=TJvgfkJpGRK129ttxY7qRYJ7ASKZA4R6-0BUA3Lk7mc,4450
 application_sdk/outputs/parquet.py,sha256=A2EnEx1zWjaXk10u3eJusmWxGxt8WR7CHXDaJgsKpq0,11040
-application_sdk/outputs/secretstore.py,sha256=AuLcM-wmzr7ZTXePLk3dunj3vhIhWkigQ-XYJo6SOZA,1306
+application_sdk/outputs/secretstore.py,sha256=JS9vUzb11leDpcMQSCnLJuE9Ww-9G3wMvCdUKBPaw9I,1342
 application_sdk/outputs/statestore.py,sha256=erviUTNyK-iLPe03i4r-G-5D3uAWdRnkW83Hnu7wI7o,3854
 application_sdk/server/__init__.py,sha256=KTqE1YPw_3WDVMWatJUuf9OOiobLM2K5SMaBrI62sCo,1568
 application_sdk/server/fastapi/__init__.py,sha256=QiIaR43QhgJaLeKs5lTghfxCAGxVxFnYAudy47akbP4,27783
@@ -129,8 +130,8 @@ application_sdk/workflows/metadata_extraction/__init__.py,sha256=jHUe_ZBQ66jx8bg
 application_sdk/workflows/metadata_extraction/sql.py,sha256=_NhszxIgmcQI6lVpjJoyJRFLwPYvJw1Dyqox_m9K2RA,11947
 application_sdk/workflows/query_extraction/__init__.py,sha256=n066_CX5RpJz6DIxGMkKS3eGSRg03ilaCtsqfJWQb7Q,117
 application_sdk/workflows/query_extraction/sql.py,sha256=kT_JQkLCRZ44ZpaC4QvPL6DxnRIIVh8gYHLqRbMI-hA,4826
-atlan_application_sdk-0.1.1rc26.dist-info/METADATA,sha256=wnjrsH3xghiBQv3BLczFc9_PpnmVhohjtlC80qHzUuo,5476
-atlan_application_sdk-0.1.1rc26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-atlan_application_sdk-0.1.1rc26.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-atlan_application_sdk-0.1.1rc26.dist-info/licenses/NOTICE,sha256=zlRshGtgfjXFcCKn_hwuILOFV9-NoKe0e2Tjy98_lrw,1044
-atlan_application_sdk-0.1.1rc26.dist-info/RECORD,,
+atlan_application_sdk-0.1.1rc28.dist-info/METADATA,sha256=wAdmNUPLPieA06duXajVG_NdweyzxwZMA60mDYShTXc,5473
+atlan_application_sdk-0.1.1rc28.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+atlan_application_sdk-0.1.1rc28.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-0.1.1rc28.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-0.1.1rc28.dist-info/RECORD,,

{atlan_application_sdk-0.1.1rc26.dist-info → atlan_application_sdk-0.1.1rc28.dist-info}/licenses/NOTICE

@@ -19,7 +19,7 @@ Apache License 2.0 Dependencies:
   - Apache Atlas (https://github.com/apache/atlas) under Apache-2.0 license
   - Elasticsearch DSL (https://github.com/elastic/elasticsearch-dsl-py) under Apache-2.0 license
 - orjson (https://github.com/ijl/orjson)
-- getdaft (https://github.com/Eventual-Inc/Daft)
+- daft (https://github.com/Eventual-Inc/Daft)
 - pyiceberg (https://github.com/apache/iceberg-python)
 - pyarrow (https://github.com/apache/arrow/tree/main/python)
 - dapr (https://github.com/dapr/python-sdk/)