atlan-application-sdk 0.1.1rc33__py3-none-any.whl → 0.1.1rc35__py3-none-any.whl

application_sdk/clients/temporal.py

@@ -42,12 +42,10 @@ from application_sdk.events.models import (
     EventTypes,
     WorkflowStates,
 )
-from application_sdk.inputs.secretstore import SecretStoreInput
-from application_sdk.inputs.statestore import StateType
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.outputs.eventstore import EventStore
-from application_sdk.outputs.secretstore import SecretStoreOutput
-from application_sdk.outputs.statestore import StateStoreOutput
+from application_sdk.services.eventstore import EventStore
+from application_sdk.services.secretstore import SecretStore
+from application_sdk.services.statestore import StateStore, StateType
 from application_sdk.workflows import WorkflowInterface
 
 logger = get_logger(__name__)
@@ -269,9 +267,7 @@ class TemporalWorkflowClient(WorkflowClient):
         self.port = port if port else WORKFLOW_PORT
         self.namespace = namespace if namespace else WORKFLOW_NAMESPACE
 
-        self.deployment_config: Dict[str, Any] = (
-            SecretStoreInput.get_deployment_secret()
-        )
+        self.deployment_config: Dict[str, Any] = SecretStore.get_deployment_secret()
         self.worker_task_queue = self.get_worker_task_queue()
         self.auth_manager = AtlanAuthClient()
 
@@ -426,7 +422,7 @@ class TemporalWorkflowClient(WorkflowClient):
         """
         if "credentials" in workflow_args:
             # remove credentials from workflow_args and add reference to credentials
-            workflow_args["credential_guid"] = await SecretStoreOutput.save_secret(
+            workflow_args["credential_guid"] = await SecretStore.save_secret(
                 workflow_args["credentials"]
             )
             del workflow_args["credentials"]
@@ -442,7 +438,7 @@ class TemporalWorkflowClient(WorkflowClient):
                 }
             )
 
-            await StateStoreOutput.save_state_object(
+            await StateStore.save_state_object(
                 id=workflow_id, value=workflow_args, type=StateType.WORKFLOWS
             )
             logger.info(f"Created workflow config with ID: {workflow_id}")

application_sdk/handlers/base.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, Optional
+
+from application_sdk.clients.base import BaseClient
+from application_sdk.handlers import HandlerInterface
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+class BaseHandler(HandlerInterface):
+    """
+    Base handler for non-SQL based applications.
+
+    This class provides a base implementation for handlers that need to interact with non-SQL data sources. It implements the HandlerInterface and provides basic functionality that can be extended by subclasses.
+
+    Attributes:
+        client (BaseClient): The client instance for connecting to the target system.
+    """
+
+    def __init__(self, client: Optional[BaseClient] = None):
+        """
+        Initialize the base handler.
+
+        Args:
+            client (BaseClient, optional): The client instance to use for connections. Defaults to BaseClient().
+        """
+        self.client = client or BaseClient()
+
+    async def load(self, credentials: Dict[str, Any]) -> None:
+        """
+        Load and initialize the handler.
+
+        This method initializes the handler and loads the client with the provided credentials.
+
+        Args:
+            credentials (Dict[str, Any]): Credentials for the client.
+        """
+        logger.info("Loading base handler")
+
+        # Load the client with credentials
+        await self.client.load(credentials=credentials)
+
+        logger.info("Base handler loaded successfully")
+
+    # The following methods are inherited from HandlerInterface and should be implemented
+    # by subclasses to handle calls from their respective FastAPI endpoints:
+    #
+    # - test_auth(**kwargs) -> bool: Called by /workflow/v1/auth endpoint
+    # - preflight_check(**kwargs) -> Any: Called by /workflow/v1/check endpoint
+    # - fetch_metadata(**kwargs) -> Any: Called by /workflow/v1/metadata endpoint

application_sdk/inputs/json.py

@@ -1,10 +1,11 @@
 import os
 from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
 
+from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.common.error_codes import IOError
 from application_sdk.inputs import Input
-from application_sdk.inputs.objectstore import ObjectStoreInput
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
 
 if TYPE_CHECKING:
     import daft
@@ -51,9 +52,10 @@ class JsonInput(Input):
                 if self.download_file_prefix is not None and not os.path.exists(
                     os.path.join(self.path, file_name)
                 ):
-                    ObjectStoreInput.download_file_from_object_store(
-                        os.path.join(self.download_file_prefix, file_name),
-                        os.path.join(self.path, file_name),
+                    destination_file_path = os.path.join(self.path, file_name)
+                    await ObjectStore.download_file(
+                        source=get_object_store_prefix(destination_file_path),
+                        destination=destination_file_path,
                     )
             except IOError as e:
                 logger.error(

application_sdk/inputs/parquet.py

@@ -2,9 +2,10 @@ import glob
 import os
 from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
 
+from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.inputs import Input
-from application_sdk.inputs.objectstore import ObjectStoreInput
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 
@@ -42,37 +43,39 @@ class ParquetInput(Input):
         self.input_prefix = input_prefix
         self.file_names = file_names
 
-    async def download_files(self, local_file_path: str) -> Optional[str]:
+    async def download_files(self, local_path: str) -> Optional[str]:
         """Read a file from the object store.
 
         Args:
-            local_file_path (str): Path to the local file in the temp directory.
+            local_path (str): Path to the local data in the temp directory.
 
         Returns:
             Optional[str]: Path to the downloaded local file.
         """
         # if the path is a directory, then check if the directory has any parquet files
         parquet_files = []
-        if os.path.isdir(local_file_path):
-            parquet_files = glob.glob(os.path.join(local_file_path, "*.parquet"))
+        if os.path.isdir(local_path):
+            parquet_files = glob.glob(os.path.join(local_path, "*.parquet"))
         else:
-            parquet_files = glob.glob(local_file_path)
+            parquet_files = glob.glob(local_path)
         if not parquet_files:
             if self.input_prefix:
                 logger.info(
-                    f"Reading file from object store: {local_file_path} from {self.input_prefix}"
+                    f"Reading file from object store: {local_path} from {self.input_prefix}"
                 )
-                if os.path.isdir(local_file_path):
-                    ObjectStoreInput.download_files_from_object_store(
-                        self.input_prefix, local_file_path
+                if os.path.isdir(local_path):
+                    await ObjectStore.download_prefix(
+                        source=get_object_store_prefix(local_path),
+                        destination=local_path,
                     )
                 else:
-                    ObjectStoreInput.download_file_from_object_store(
-                        self.input_prefix, local_file_path
+                    await ObjectStore.download_file(
+                        source=get_object_store_prefix(local_path),
+                        destination=local_path,
                     )
             else:
                 raise ValueError(
-                    f"No parquet files found in {local_file_path} and no input prefix provided"
+                    f"No parquet files found in {local_path} and no input prefix provided"
                 )
 
     async def get_dataframe(self) -> "pd.DataFrame":

application_sdk/outputs/__init__.py

@@ -22,9 +22,10 @@ import orjson
 from temporalio import activity
 
 from application_sdk.activities.common.models import ActivityStatistics
+from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.common.dataframe_utils import is_empty_dataframe
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.outputs.objectstore import ObjectStoreOutput
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 activity.logger = logger
@@ -223,9 +224,11 @@ class Output(ABC):
             with open(output_file_name, "w") as f:
                 f.write(orjson.dumps(statistics).decode("utf-8"))
 
+            destination_file_path = get_object_store_prefix(output_file_name)
             # Push the file to the object store
-            await ObjectStoreOutput.push_file_to_object_store(
-                self.output_prefix, output_file_name
+            await ObjectStore.upload_file(
+                source=output_file_name,
+                destination=destination_file_path,
             )
             return statistics
         except Exception as e:

application_sdk/outputs/json.py

@@ -5,10 +5,11 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 import orjson
 from temporalio import activity
 
+from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
 from application_sdk.outputs import Output
-from application_sdk.outputs.objectstore import ObjectStoreOutput
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 activity.logger = logger
@@ -285,9 +286,10 @@ class JsonOutput(Output):
                 description="Number of records written to JSON files from daft DataFrame",
             )
 
-            # Push the file to the object store
-            await ObjectStoreOutput.push_files_to_object_store(
-                self.output_prefix, self.output_path
+            # Push files to the object store
+            await ObjectStore.upload_prefix(
+                source=self.output_path,
+                destination=get_object_store_prefix(self.output_path),
             )
 
         except Exception as e:
@@ -344,8 +346,9 @@ class JsonOutput(Output):
                 )
 
                 # Push the file to the object store
-                await ObjectStoreOutput.push_file_to_object_store(
-                    self.output_prefix, output_file_name
+                await ObjectStore.upload_file(
+                    source=output_file_name,
+                    destination=get_object_store_prefix(output_file_name),
                 )
 
             self.buffer.clear()

application_sdk/outputs/parquet.py

@@ -3,10 +3,11 @@ from typing import TYPE_CHECKING, Literal, Optional
 
 from temporalio import activity
 
+from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
 from application_sdk.outputs import Output
-from application_sdk.outputs.objectstore import ObjectStoreOutput
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 activity.logger = logger
@@ -159,7 +160,10 @@ class ParquetOutput(Output):
             )
 
             # Upload the file to object store
-            await self.upload_file(file_path)
+            await ObjectStore.upload_file(
+                source=file_path,
+                destination=get_object_store_prefix(file_path),
+            )
         except Exception as e:
             # Record metrics for failed write
             self.metrics.record_metric(
@@ -218,7 +222,10 @@ class ParquetOutput(Output):
             )
 
             # Upload the file to object store
-            await self.upload_file(file_path)
+            await ObjectStore.upload_file(
+                source=file_path,
+                destination=get_object_store_prefix(file_path),
+            )
         except Exception as e:
             # Record metrics for failed write
             self.metrics.record_metric(
@@ -231,39 +238,6 @@ class ParquetOutput(Output):
             logger.error(f"Error writing daft dataframe to parquet: {str(e)}")
             raise
 
-    async def upload_file(self, local_file_path: str) -> None:
-        """Upload a file to the object store.
-
-        Args:
-            local_file_path (str): Path to the local file to upload.
-        """
-        try:
-            if os.path.isdir(local_file_path):
-                logger.info(
-                    f"Uploading files: {local_file_path} to {self.output_prefix}"
-                )
-                await ObjectStoreOutput.push_files_to_object_store(
-                    self.output_prefix, local_file_path
-                )
-            else:
-                logger.info(
-                    f"Uploading file: {local_file_path} to {self.output_prefix}"
-                )
-                await ObjectStoreOutput.push_file_to_object_store(
-                    self.output_prefix, local_file_path
-                )
-        except Exception as e:
-            # Record metrics for failed upload
-            self.metrics.record_metric(
-                name="parquet_upload_errors",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={"error": str(e)},
-                description="Number of errors while uploading Parquet files to object store",
-            )
-            logger.error(f"Error uploading file to object store: {str(e)}")
-            raise e
-
     def get_full_path(self) -> str:
         """Get the full path of the output file.
 

application_sdk/server/fastapi/__init__.py

@@ -25,11 +25,9 @@ from application_sdk.constants import (
 )
 from application_sdk.docgen import AtlanDocsGenerator
 from application_sdk.handlers import HandlerInterface
-from application_sdk.inputs.statestore import StateStoreInput, StateType
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
 from application_sdk.observability.observability import DuckDBUI
-from application_sdk.outputs.statestore import StateStoreOutput
 from application_sdk.server import ServerInterface
 from application_sdk.server.fastapi.middleware.logmiddleware import LogMiddleware
 from application_sdk.server.fastapi.middleware.metrics import MetricsMiddleware
@@ -53,6 +51,7 @@ from application_sdk.server.fastapi.models import (
 )
 from application_sdk.server.fastapi.routers.server import get_server_router
 from application_sdk.server.fastapi.utils import internal_server_error_handler
+from application_sdk.services.statestore import StateStore, StateType
 from application_sdk.workflows import WorkflowInterface
 
 logger = get_logger(__name__)
@@ -588,7 +587,7 @@ class APIServer(ServerInterface):
             )
             raise e
 
-    def get_workflow_config(
+    async def get_workflow_config(
         self, config_id: str, type: str = "workflows"
     ) -> WorkflowConfigResponse:
         """Retrieve workflow configuration by ID.
@@ -603,7 +602,7 @@ class APIServer(ServerInterface):
         if not StateType.is_member(type):
             raise ValueError(f"Invalid type {type} for state store")
 
-        config = StateStoreInput.get_state(config_id, StateType(type))
+        config = await StateStore.get_state(config_id, StateType(type))
         return WorkflowConfigResponse(
             success=True,
             message="Workflow configuration fetched successfully",
@@ -680,7 +679,7 @@ class APIServer(ServerInterface):
         if not StateType.is_member(type):
             raise ValueError(f"Invalid type {type} for state store")
 
-        config = await StateStoreOutput.save_state_object(
+        config = await StateStore.save_state_object(
             id=config_id, value=body.model_dump(), type=StateType(type)
         )
         return WorkflowConfigResponse(

application_sdk/server/fastapi/models.py

@@ -34,7 +34,7 @@ class FetchMetadataRequest(RootModel[Dict[str, Any]]):
 
 class FetchMetadataResponse(BaseModel):
     success: bool
-    data: List[Dict[str, str]]
+    data: Any
 
 
 class PreflightCheckRequest(BaseModel):

application_sdk/services/__init__.py

@@ -0,0 +1,18 @@
+"""Services module for the application SDK."""
+
+from .atlan_storage import AtlanStorage, MigrationSummary
+from .eventstore import EventStore
+from .objectstore import ObjectStore
+from .secretstore import SecretStore
+from .statestore import StateStore, StateType, build_state_store_path
+
+__all__ = [
+    "AtlanStorage",
+    "EventStore",
+    "MigrationSummary",
+    "ObjectStore",
+    "SecretStore",
+    "StateStore",
+    "StateType",
+    "build_state_store_path",
+]

application_sdk/{outputs → services}/atlan_storage.py

@@ -1,4 +1,12 @@
-"""Atlan storage interface for upload operations and migration from objectstore."""
+"""Atlan storage service for upload operations and migration from object store.
+
+This module provides the AtlanStorage service for handling data migration between
+local object storage and Atlan's upstream storage system. It's specifically designed
+for the bucket cloning strategy used in customer-deployed applications.
+
+The service supports parallel file migration with comprehensive error handling and
+detailed reporting through the MigrationSummary model.
+"""
 
 import asyncio
 from typing import Dict, List
@@ -11,8 +19,8 @@ from application_sdk.constants import (
     DEPLOYMENT_OBJECT_STORE_NAME,
     UPSTREAM_OBJECT_STORE_NAME,
 )
-from application_sdk.inputs.objectstore import ObjectStoreInput
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 activity.logger = logger
@@ -43,27 +51,37 @@ class MigrationSummary(BaseModel):
     destination: str = UPSTREAM_OBJECT_STORE_NAME
 
 
-# keeping any logic related to operations on atlan storage within this file.
-class AtlanStorageOutput:
+class AtlanStorage:
     """Handles upload operations to Atlan storage and migration from objectstore."""
 
     OBJECT_CREATE_OPERATION = "create"
 
     @classmethod
     async def _migrate_single_file(cls, file_path: str) -> tuple[str, bool, str]:
-        """
-        Migrate a single file from objectstore to Atlan storage.
+        """Migrate a single file from object store to Atlan storage.
+
+        This internal method handles the migration of a single file, including
+        error handling and logging. It's designed to be called concurrently
+        for multiple files.
 
         Args:
-            file_path (str): The path of the file to migrate
+            file_path (str): The path of the file to migrate in the object store.
 
         Returns:
-            tuple[str, bool, str]: (file_path, success, error_message)
+            tuple[str, bool, str]: A tuple containing:
+                - file_path: The path of the file that was processed
+                - success: Boolean indicating if migration was successful
+                - error_message: Error details if migration failed, empty string if successful
+
+        Note:
+            This method is internal and should not be called directly. Use
+            migrate_from_objectstore_to_atlan() instead for proper coordination
+            and error handling.
         """
         try:
             # Get file data from objectstore
-            file_data = ObjectStoreInput.get_file_data(
-                file_path, object_store_name=DEPLOYMENT_OBJECT_STORE_NAME
+            file_data = await ObjectStore.get_content(
+                file_path, store_name=DEPLOYMENT_OBJECT_STORE_NAME
             )
 
             with DaprClient() as client:
@@ -91,14 +109,44 @@ class AtlanStorageOutput:
     async def migrate_from_objectstore_to_atlan(
         cls, prefix: str = ""
     ) -> MigrationSummary:
-        """
-        Migrate all files from objectstore to Atlan storage under a given prefix.
+        """Migrate all files from object store to Atlan storage under a given prefix.
+
+        This method performs a parallel migration of files from the local object store
+        to Atlan's upstream storage system. It provides comprehensive error handling
+        and detailed reporting of the migration process.
 
         Args:
-            prefix (str): The prefix to filter which files to migrate. Empty string migrates all files.
+            prefix (str, optional): The prefix to filter which files to migrate.
+                Empty string migrates all files. Defaults to "".
 
         Returns:
-            MigrationSummary: Migration summary with counts and any failures
+            MigrationSummary: Comprehensive migration summary including:
+                - total_files: Number of files found for migration
+                - migrated_files: Number successfully migrated
+                - failed_migrations: Number that failed to migrate
+                - failures: List of failure details with file paths and errors
+                - prefix: The prefix used for filtering
+                - source/destination: Storage system identifiers
+
+        Raises:
+            Exception: If there's a critical error during the migration process.
+
+        Examples:
+            >>> # Migrate all files
+            >>> summary = await AtlanStorage.migrate_from_objectstore_to_atlan()
+            >>> print(f"Success rate: {summary.migrated_files/summary.total_files*100:.1f}%")
+
+            >>> # Migrate specific dataset
+            >>> summary = await AtlanStorage.migrate_from_objectstore_to_atlan(
+            ...     prefix="processed_data/2024/"
+            ... )
+            >>> if summary.total_files == 0:
+            ...     print("No files found with the specified prefix")
+            >>> elif summary.failed_migrations == 0:
+            ...     print(f"Successfully migrated all {summary.total_files} files")
+            >>> else:
+            ...     print(f"Migration completed with {summary.failed_migrations} failures")
+            ...     # Handle failures...
         """
         try:
             logger.info(
@@ -106,8 +154,8 @@ class AtlanStorageOutput:
             )
 
             # Get list of all files to migrate from objectstore
-            files_to_migrate = ObjectStoreInput.list_all_files(
-                prefix, object_store_name=DEPLOYMENT_OBJECT_STORE_NAME
+            files_to_migrate = await ObjectStore.list_files(
+                prefix, store_name=DEPLOYMENT_OBJECT_STORE_NAME
             )
 
             total_files = len(files_to_migrate)

application_sdk/{outputs → services}/eventstore.py

@@ -1,4 +1,4 @@
-"""Event store module for handling application events.
+"""Unified event store service for handling application events.
 
 This module provides the EventStore class for publishing application events
 to a pub/sub system with automatic fallback to HTTP binding.
@@ -10,7 +10,6 @@ from datetime import datetime
 from dapr import clients
 from temporalio import activity, workflow
 
-from application_sdk.clients.atlan_auth import AtlanAuthClient
 from application_sdk.common.dapr_utils import is_component_registered
 from application_sdk.constants import (
     APPLICATION_NAME,
@@ -25,17 +24,39 @@ activity.logger = logger
 
 
 class EventStore:
-    """Event store for publishing application events.
+    """Unified event store service for publishing application events.
 
     This class provides functionality to publish events to a pub/sub system.
     """
 
     @classmethod
     def enrich_event_metadata(cls, event: Event):
-        """Enrich the event metadata with the workflow and activity information.
+        """Enrich the event metadata with workflow and activity context information.
+
+        This method automatically populates event metadata with context from the current
+        Temporal workflow and activity execution, including IDs, types, and execution state.
 
         Args:
-            event (Event): Event data.
+            event (Event): Event data to enrich with metadata.
+
+        Returns:
+            Event: The same event instance with enriched metadata.
+
+        Note:
+            This method safely handles cases where the code is not running within
+            a Temporal workflow or activity context.
+
+        Examples:
+            >>> from application_sdk.events.models import Event
+
+            >>> # Create basic event
+            >>> event = Event(event_type="data.processed", data={"count": 100})
+
+            >>> # Enrich with current context (if available)
+            >>> enriched = EventStore.enrich_event_metadata(event)
+            >>> print(f"Workflow ID: {enriched.metadata.workflow_id}")
+            >>> print(f"Activity: {enriched.metadata.activity_type}")
+            >>> print(f"Timestamp: {enriched.metadata.created_timestamp}")
         """
         if not event.metadata:
             event.metadata = EventMetadata()
@@ -70,10 +91,49 @@ class EventStore:
 
     @classmethod
     async def publish_event(cls, event: Event):
-        """Publish event with automatic fallback between pub/sub and HTTP binding.
+        """Publish event with automatic metadata enrichment and authentication.
+
+        This method handles the complete event publishing flow including metadata
+        enrichment, authentication header injection, and component availability validation.
+        It automatically falls back gracefully if the event store component is not available.
 
         Args:
             event (Event): Event data to publish.
+
+        Note:
+            The method will silently skip publishing if the event store component
+            is not registered, allowing applications to run without event publishing
+            capability.
+
+        Raises:
+            Exception: If there's an error during event publishing (logged but not re-raised).
+
+        Examples:
+            >>> from application_sdk.events.models import Event
+
+            >>> # Publish workflow status event
+            >>> status_event = Event(
+            ...     event_type="workflow.status_changed",
+            ...     data={
+            ...         "workflow_id": "wf-123",
+            ...         "old_status": "running",
+            ...         "new_status": "completed",
+            ...         "duration_seconds": 1800
+            ...     }
+            ... )
+            >>> await EventStore.publish_event(status_event)
+
+            >>> # Publish data processing event
+            >>> processing_event = Event(
+            ...     event_type="data.batch_processed",
+            ...     data={
+            ...         "batch_id": "batch-456",
+            ...         "records_processed": 10000,
+            ...         "success_count": 9995,
+            ...         "error_count": 5
+            ...     }
+            ... )
+            >>> await EventStore.publish_event(processing_event)
         """
         if not is_component_registered(EVENT_STORE_NAME):
             logger.warning(
@@ -89,6 +149,8 @@ class EventStore:
             binding_metadata = {"content-type": "application/json"}
 
             # Add auth token - HTTP bindings will use it, others will ignore it
+            from application_sdk.clients.atlan_auth import AtlanAuthClient
+
             auth_client = AtlanAuthClient()
             binding_metadata.update(await auth_client.get_authenticated_headers())