atlan-application-sdk 0.1.1rc34__py3-none-any.whl → 0.1.1rc36__py3-none-any.whl

application_sdk/interceptors/events.py

@@ -0,0 +1,193 @@
+from datetime import timedelta
+from typing import Any, Optional, Type
+
+from temporalio import activity, workflow
+from temporalio.common import RetryPolicy
+from temporalio.worker import (
+    ActivityInboundInterceptor,
+    ExecuteActivityInput,
+    ExecuteWorkflowInput,
+    Interceptor,
+    WorkflowInboundInterceptor,
+    WorkflowInterceptorClassInput,
+)
+
+from application_sdk.events.models import (
+    ApplicationEventNames,
+    Event,
+    EventMetadata,
+    EventTypes,
+    WorkflowStates,
+)
+from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.eventstore import EventStore
+
+logger = get_logger(__name__)
+
+TEMPORAL_NOT_FOUND_FAILURE = (
+    "type.googleapis.com/temporal.api.errordetails.v1.NotFoundFailure"
+)
+
+
+# Activity for publishing events (runs outside sandbox)
+@activity.defn
+async def publish_event(event_data: dict) -> None:
+    """Activity to publish events outside the workflow sandbox.
+
+    Args:
+        event_data (dict): Event data to publish containing event_type, event_name,
+                          metadata, and data fields.
+    """
+    try:
+        event = Event(**event_data)
+        await EventStore.publish_event(event)
+        activity.logger.info(f"Published event: {event_data.get('event_name','')}")
+    except Exception as e:
+        activity.logger.error(f"Failed to publish event: {e}")
+        raise
+
+
+class EventActivityInboundInterceptor(ActivityInboundInterceptor):
+    """Interceptor for tracking activity execution events.
+
+    This interceptor captures the start and end of activity executions,
+    creating events that can be used for monitoring and tracking.
+    Activities run outside the sandbox so they can directly call EventStore.
+    """
+
+    async def execute_activity(self, input: ExecuteActivityInput) -> Any:
+        """Execute an activity with event tracking.
+
+        Args:
+            input (ExecuteActivityInput): The activity execution input.
+
+        Returns:
+            Any: The result of the activity execution.
+        """
+        # Extract activity information for tracking
+
+        start_event = Event(
+            event_type=EventTypes.APPLICATION_EVENT.value,
+            event_name=ApplicationEventNames.ACTIVITY_START.value,
+            data={},
+        )
+        await EventStore.publish_event(start_event)
+
+        output = None
+        try:
+            output = await super().execute_activity(input)
+        except Exception:
+            raise
+        finally:
+            end_event = Event(
+                event_type=EventTypes.APPLICATION_EVENT.value,
+                event_name=ApplicationEventNames.ACTIVITY_END.value,
+                data={},
+            )
+            await EventStore.publish_event(end_event)
+
+        return output
+
+
+class EventWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+    """Interceptor for tracking workflow execution events.
+
+    This interceptor captures the start and end of workflow executions,
+    creating events that can be used for monitoring and tracking.
+    Uses activities to publish events to avoid sandbox restrictions.
+    """
+
+    async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
+        """Execute a workflow with event tracking.
+
+        Args:
+            input (ExecuteWorkflowInput): The workflow execution input.
+
+        Returns:
+            Any: The result of the workflow execution.
+        """
+
+        # Publish workflow start event via activity
+        try:
+            await workflow.execute_activity(
+                publish_event,
+                {
+                    "metadata": EventMetadata(
+                        workflow_state=WorkflowStates.RUNNING.value
+                    ),
+                    "event_type": EventTypes.APPLICATION_EVENT.value,
+                    "event_name": ApplicationEventNames.WORKFLOW_START.value,
+                    "data": {},
+                },
+                schedule_to_close_timeout=timedelta(seconds=30),
+                retry_policy=RetryPolicy(maximum_attempts=3),
+            )
+        except Exception as e:
+            workflow.logger.warning(f"Failed to publish workflow start event: {e}")
+            # Don't fail the workflow if event publishing fails
+
+        output = None
+        workflow_state = WorkflowStates.FAILED.value  # Default to failed
+
+        try:
+            output = await super().execute_workflow(input)
+            workflow_state = (
+                WorkflowStates.COMPLETED.value
+            )  # Update to completed on success
+        except Exception:
+            workflow_state = WorkflowStates.FAILED.value  # Keep as failed
+            raise
+        finally:
+            # Always publish workflow end event
+            try:
+                await workflow.execute_activity(
+                    publish_event,
+                    {
+                        "metadata": EventMetadata(workflow_state=workflow_state),
+                        "event_type": EventTypes.APPLICATION_EVENT.value,
+                        "event_name": ApplicationEventNames.WORKFLOW_END.value,
+                        "data": {},
+                    },
+                    schedule_to_close_timeout=timedelta(seconds=30),
+                    retry_policy=RetryPolicy(maximum_attempts=3),
+                )
+            except Exception as publish_error:
+                workflow.logger.warning(
+                    f"Failed to publish workflow end event: {publish_error}"
+                )
+
+        return output
+
+
+class EventInterceptor(Interceptor):
+    """Temporal interceptor for event tracking.
+
+    This interceptor provides event tracking capabilities for both
+    workflow and activity executions.
+    """
+
+    def intercept_activity(
+        self, next: ActivityInboundInterceptor
+    ) -> ActivityInboundInterceptor:
+        """Intercept activity executions.
+
+        Args:
+            next (ActivityInboundInterceptor): The next interceptor in the chain.
+
+        Returns:
+            ActivityInboundInterceptor: The activity interceptor.
+        """
+        return EventActivityInboundInterceptor(super().intercept_activity(next))
+
+    def workflow_interceptor_class(
+        self, input: WorkflowInterceptorClassInput
+    ) -> Optional[Type[WorkflowInboundInterceptor]]:
+        """Get the workflow interceptor class.
+
+        Args:
+            input (WorkflowInterceptorClassInput): The interceptor input.
+
+        Returns:
+            Optional[Type[WorkflowInboundInterceptor]]: The workflow interceptor class.
+        """
+        return EventWorkflowInboundInterceptor

application_sdk/interceptors/lock.py

@@ -0,0 +1,139 @@
+"""Redis lock interceptor for Temporal workflows.
+
+Manages distributed locks for activities decorated with @needs_lock using
+separate lock acquisition and release activities to avoid workflow deadlocks.
+"""
+
+from datetime import timedelta
+from typing import Any, Dict, Optional, Type
+
+from temporalio import workflow
+from temporalio.common import RetryPolicy
+from temporalio.worker import (
+    Interceptor,
+    StartActivityInput,
+    WorkflowInboundInterceptor,
+    WorkflowInterceptorClassInput,
+    WorkflowOutboundInterceptor,
+)
+
+from application_sdk.common.error_codes import WorkflowError
+from application_sdk.constants import (
+    APPLICATION_NAME,
+    IS_LOCKING_DISABLED,
+    LOCK_METADATA_KEY,
+)
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+class RedisLockInterceptor(Interceptor):
+    """Main interceptor class for Redis distributed locking."""
+
+    def __init__(self, activities: Dict[str, Any]):
+        """Initialize Redis lock interceptor.
+
+        Args:
+            activities: Dictionary mapping activity names to activity functions
+        """
+        self.activities = activities
+
+    def workflow_interceptor_class(
+        self, input: WorkflowInterceptorClassInput
+    ) -> Optional[Type[WorkflowInboundInterceptor]]:
+        activities = self.activities
+
+        class RedisLockWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+            """Inbound interceptor that manages Redis locks for activities."""
+
+            def init(self, outbound: WorkflowOutboundInterceptor) -> None:
+                """Initialize with Redis lock outbound interceptor."""
+                lock_outbound = RedisLockOutboundInterceptor(outbound, activities)
+                super().init(lock_outbound)
+
+        return RedisLockWorkflowInboundInterceptor
+
+
+class RedisLockOutboundInterceptor(WorkflowOutboundInterceptor):
+    """Outbound interceptor that acquires Redis locks before activity execution."""
+
+    def __init__(self, next: WorkflowOutboundInterceptor, activities: Dict[str, Any]):
+        super().__init__(next)
+        self.activities = activities
+
+    async def start_activity(  # type: ignore[override]
+        self, input: StartActivityInput
+    ) -> workflow.ActivityHandle[Any]:
+        """Start activity with distributed lock if required."""
+
+        # Check if activity needs locking
+        activity_fn = self.activities.get(input.activity)
+        if (
+            not activity_fn
+            or not hasattr(activity_fn, LOCK_METADATA_KEY)
+            or IS_LOCKING_DISABLED
+        ):
+            return await self.next.start_activity(input)
+
+        lock_config = getattr(activity_fn, LOCK_METADATA_KEY)
+        lock_name = lock_config.get("lock_name", input.activity)
+        max_locks = lock_config.get("max_locks", 5)
+        if not input.schedule_to_close_timeout:
+            logger.error(
+                f"Activity '{input.activity}' with @needs_lock decorator requires schedule_to_close_timeout"
+            )
+            raise WorkflowError(
+                f"{WorkflowError.WORKFLOW_CONFIG_ERROR}: Activity '{input.activity}' with @needs_lock decorator must be called with schedule_to_close_timeout parameter. "
+                f"Example: workflow.execute_activity('{input.activity}', schedule_to_close_timeout=timedelta(minutes=10))"
+            )
+        ttl_seconds = int(input.schedule_to_close_timeout.total_seconds())
+
+        # Orchestrate lock acquisition -> business activity -> lock release
+        return await self._execute_with_lock_orchestration(
+            input, lock_name, max_locks, ttl_seconds
+        )
+
+    async def _execute_with_lock_orchestration(
+        self,
+        input: StartActivityInput,
+        lock_name: str,
+        max_locks: int,
+        ttl_seconds: int,
+    ) -> workflow.ActivityHandle[Any]:
+        """Execute activity with distributed lock orchestration."""
+        owner_id = f"{APPLICATION_NAME}:{workflow.info().run_id}"
+        lock_result = None
+
+        try:
+            # Step 1: Acquire lock via dedicated activity (can take >2s safely)
+            start_to_close_timeout = workflow.info().execution_timeout
+            lock_result = await workflow.execute_activity(
+                "acquire_distributed_lock",
+                args=[lock_name, max_locks, ttl_seconds, owner_id],
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=RetryPolicy(maximum_attempts=1),
+            )
+
+            logger.debug(f"Lock acquired: {lock_result}, executing {input.activity}")
+
+            # Step 2: Execute the business activity and return its handle
+            return await self.next.start_activity(input)
+
+        finally:
+            # Step 3: Release lock (fire-and-forget with short timeout)
+            if lock_result is not None:
+                try:
+                    await workflow.execute_local_activity(
+                        "release_distributed_lock",
+                        args=[lock_result["resource_id"], lock_result["owner_id"]],
+                        start_to_close_timeout=timedelta(seconds=5),
+                        retry_policy=RetryPolicy(maximum_attempts=1),
+                    )
+                    logger.debug(f"Lock released: {lock_result['resource_id']}")
+                except Exception as e:
+                    # Silent failure - TTL will handle cleanup
+                    logger.warning(
+                        f"Lock release failed for {lock_result['resource_id']}: {e}. "
+                        f"TTL will handle cleanup."
+                    )

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/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)