atlan-application-sdk 0.1.1rc35__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,15 +26,9 @@ 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.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.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
@@ -55,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.
 
@@ -537,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
@@ -552,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/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