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

application_sdk/clients/temporal.py

@@ -1,22 +1,12 @@
 import asyncio
 import uuid
 from concurrent.futures import ThreadPoolExecutor
-from datetime import timedelta
 from typing import Any, Dict, Optional, Sequence, Type
 
 from temporalio import activity, workflow
 from temporalio.client import Client, WorkflowExecutionStatus, WorkflowFailureError
-from temporalio.common import RetryPolicy
 from temporalio.types import CallableType, ClassType
-from temporalio.worker import (
-    ActivityInboundInterceptor,
-    ExecuteActivityInput,
-    ExecuteWorkflowInput,
-    Interceptor,
-    Worker,
-    WorkflowInboundInterceptor,
-    WorkflowInterceptorClassInput,
-)
+from temporalio.worker import Worker
 from temporalio.worker.workflow_sandbox import (
     SandboxedWorkflowRunner,
     SandboxRestrictions,
@@ -28,6 +18,7 @@ from application_sdk.constants import (
     APPLICATION_NAME,
     DEPLOYMENT_NAME,
     DEPLOYMENT_NAME_KEY,
+    IS_LOCKING_DISABLED,
     MAX_CONCURRENT_ACTIVITIES,
     WORKFLOW_HOST,
     WORKFLOW_MAX_TIMEOUT_HOURS,
@@ -35,19 +26,11 @@ from application_sdk.constants import (
     WORKFLOW_PORT,
     WORKFLOW_TLS_ENABLED_KEY,
 )
-from application_sdk.events.models import (
-    ApplicationEventNames,
-    Event,
-    EventMetadata,
-    EventTypes,
-    WorkflowStates,
-)
-from application_sdk.inputs.secretstore import SecretStoreInput
-from application_sdk.inputs.statestore import StateType
+from application_sdk.interceptors.events import EventInterceptor, publish_event
+from application_sdk.interceptors.lock import RedisLockInterceptor
 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.secretstore import SecretStore
+from application_sdk.services.statestore import StateStore, StateType
 from application_sdk.workflows import WorkflowInterface
 
 logger = get_logger(__name__)
@@ -57,170 +40,6 @@ TEMPORAL_NOT_FOUND_FAILURE = (
 )
 
 
-# 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
-
-
 class TemporalWorkflowClient(WorkflowClient):
     """Temporal-specific implementation of WorkflowClient with simple token refresh.
 
@@ -269,9 +88,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 +243,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 +259,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}")
@@ -541,14 +358,34 @@ class TemporalWorkflowClient(WorkflowClient):
                 f"Started token refresh loop with dynamic interval (initial: {self._token_refresh_interval}s)"
             )
 
-        # Add the publish_event to the activities list
-        extended_activities = list(activities) + [publish_event]
+        # Start with provided activities and add system activities
+        final_activities = list(activities) + [publish_event]
+
+        # Add lock management activities if needed
+        if not IS_LOCKING_DISABLED:
+            from application_sdk.activities.lock_management import (
+                acquire_distributed_lock,
+                release_distributed_lock,
+            )
+
+            final_activities.extend(
+                [
+                    acquire_distributed_lock,
+                    release_distributed_lock,
+                ]
+            )
+            logger.info(
+                "Auto-registered lock management activities for @needs_lock decorated activities"
+            )
+
+        # Create activities lookup dict for interceptors
+        activities_dict = {getattr(a, "__name__", str(a)): a for a in final_activities}
 
         return Worker(
             self.client,
             task_queue=self.worker_task_queue,
             workflows=workflow_classes,
-            activities=extended_activities,  # Use extended activities list
+            activities=final_activities,
             workflow_runner=SandboxedWorkflowRunner(
                 restrictions=SandboxRestrictions.default.with_passthrough_modules(
                     *passthrough_modules
@@ -556,7 +393,10 @@ class TemporalWorkflowClient(WorkflowClient):
             ),
             max_concurrent_activities=max_concurrent_activities,
             activity_executor=activity_executor,
-            interceptors=[EventInterceptor()],
+            interceptors=[
+                EventInterceptor(),
+                RedisLockInterceptor(activities_dict),
+            ],
         )
 
     async def get_workflow_run_status(

application_sdk/common/error_codes.py

@@ -81,6 +81,18 @@ class ClientError(AtlanError):
     AUTH_CONFIG_ERROR = ErrorCode(
         ErrorComponent.CLIENT, "400", "00", "Authentication configuration error"
     )
+    REDIS_CONNECTION_ERROR = ErrorCode(
+        ErrorComponent.CLIENT, "503", "00", "Redis connection failed"
+    )
+    REDIS_TIMEOUT_ERROR = ErrorCode(
+        ErrorComponent.CLIENT, "408", "00", "Redis operation timeout"
+    )
+    REDIS_AUTH_ERROR = ErrorCode(
+        ErrorComponent.CLIENT, "401", "05", "Redis authentication failed"
+    )
+    REDIS_PROTOCOL_ERROR = ErrorCode(
+        ErrorComponent.CLIENT, "502", "00", "Redis protocol error"
+    )
 
 
 class ApiError(AtlanError):
@@ -174,7 +186,7 @@ class IOError(AtlanError):
     INPUT_PROCESSING_ERROR = ErrorCode(
         ErrorComponent.IO, "500", "01", "Input processing error"
     )
-    SQL_QUERY_ERROR = ErrorCode(ErrorComponent.IO, "400", "00", "SQL query error")
+    SQL_QUERY_ERROR = ErrorCode(ErrorComponent.IO, "400", "01", "SQL query error")
     SQL_QUERY_BATCH_ERROR = ErrorCode(
         ErrorComponent.IO, "500", "02", "SQL query batch error"
     )
@@ -268,10 +280,10 @@ class CommonError(AtlanError):
         ErrorComponent.COMMON, "400", "01", "Query preparation error"
     )
     FILTER_PREPARATION_ERROR = ErrorCode(
-        ErrorComponent.COMMON, "400", "00", "Filter preparation error"
+        ErrorComponent.COMMON, "400", "02", "Filter preparation error"
     )
     CREDENTIALS_PARSE_ERROR = ErrorCode(
-        ErrorComponent.COMMON, "400", "02", "Credentials parse error"
+        ErrorComponent.COMMON, "400", "03", "Credentials parse error"
     )
     CREDENTIALS_RESOLUTION_ERROR = ErrorCode(
         ErrorComponent.COMMON, "401", "03", "Credentials resolution error"
@@ -367,3 +379,12 @@ class ActivityError(AtlanError):
     ATLAN_UPLOAD_ERROR = ErrorCode(
         ErrorComponent.ACTIVITY, "500", "08", "Atlan upload error"
     )
+    LOCK_ACQUISITION_ERROR = ErrorCode(
+        ErrorComponent.ACTIVITY, "503", "01", "Distributed lock acquisition error"
+    )
+    LOCK_RELEASE_ERROR = ErrorCode(
+        ErrorComponent.ACTIVITY, "500", "09", "Distributed lock release error"
+    )
+    LOCK_TIMEOUT_ERROR = ErrorCode(
+        ErrorComponent.ACTIVITY, "408", "00", "Lock acquisition timeout"
+    )

application_sdk/constants.py

@@ -146,7 +146,6 @@ DEPLOYMENT_SECRET_STORE_NAME = os.getenv(
     "DEPLOYMENT_SECRET_STORE_NAME", "deployment-secret-store"
 )
 
-
 # Logger Constants
 #: Log level for the application (DEBUG, INFO, WARNING, ERROR, CRITICAL)
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
@@ -230,3 +229,21 @@ 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")
+# Lock Configuration
+LOCK_METADATA_KEY = "__lock_metadata__"
+
+# Redis Lock Configuration
+#: Redis host for direct connection (when not using Sentinel)
+REDIS_HOST = os.getenv("REDIS_HOST", "")
+#: Redis port for direct connection (when not using Sentinel)
+REDIS_PORT = os.getenv("REDIS_PORT", "")
+#: Redis password (required for authenticated Redis instances)
+REDIS_PASSWORD = os.getenv("REDIS_PASSWORD")
+#: Redis Sentinel service name (default: mymaster)
+REDIS_SENTINEL_SERVICE_NAME = os.getenv("REDIS_SENTINEL_SERVICE_NAME", "mymaster")
+#: Redis Sentinel hosts (comma-separated host:port pairs)
+REDIS_SENTINEL_HOSTS = os.getenv("REDIS_SENTINEL_HOSTS", "")
+#: Whether to enable strict locking
+IS_LOCKING_DISABLED = os.getenv("IS_LOCKING_DISABLED", "true").lower() == "true"
+#: Retry interval for lock acquisition
+LOCK_RETRY_INTERVAL = int(os.getenv("LOCK_RETRY_INTERVAL", "5"))

application_sdk/decorators/locks.py

@@ -0,0 +1,42 @@
+from typing import Any, Callable, Optional
+
+from application_sdk.constants import LOCK_METADATA_KEY
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+def needs_lock(max_locks: int = 5, lock_name: Optional[str] = None):
+    """Decorator to mark activities that require distributed locking.
+
+    This decorator attaches lock configuration directly to the activity
+    definition that will be used by the workflow interceptor to acquire
+    locks before executing activities.
+
+    Note:
+        Activities decorated with ``needs_lock`` must be called with
+        ``schedule_to_close_timeout`` to ensure proper lock TTL calculation
+        that covers retries.
+
+    Args:
+        max_locks (int): Maximum number of concurrent locks allowed.
+        lock_name (str | None): Optional custom name for the lock (defaults to activity name).
+
+    Raises:
+        WorkflowError: If activity is called without ``schedule_to_close_timeout``.
+    """
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        # Store lock metadata directly on the function object
+        metadata = {
+            "is_needs_lock": True,
+            "max_locks": max_locks,
+            "lock_name": lock_name or func.__name__,
+        }
+
+        # Attach metadata to the function
+        setattr(func, LOCK_METADATA_KEY, metadata)
+
+        return func
+
+    return decorator

application_sdk/handlers/base.py

@@ -44,7 +44,24 @@ class BaseHandler(HandlerInterface):
 
     # 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
+
+    async def test_auth(self, **kwargs: Any) -> bool:
+        """
+        Test the authentication of the handler.
+        """
+        raise NotImplementedError("test_auth is not implemented")
+
+    async def preflight_check(self, **kwargs: Any) -> Any:
+        """
+        Check the preflight of the handler.
+        """
+        raise NotImplementedError("preflight_check is not implemented")
+
+    async def fetch_metadata(self, **kwargs: Any) -> Any:
+        """
+        Fetch the metadata of the handler.
+        """
+        raise NotImplementedError("fetch_metadata is not implemented")

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":