atlan-application-sdk 0.1.1rc39__py3-none-any.whl → 0.1.1rc41__py3-none-any.whl

application_sdk/interceptors/cleanup.py

@@ -0,0 +1,171 @@
+import os
+import shutil
+from datetime import timedelta
+from typing import Any, Dict, List, Optional, Type
+
+from pydantic import BaseModel
+from temporalio import activity, workflow
+from temporalio.common import RetryPolicy
+from temporalio.worker import (
+    ExecuteWorkflowInput,
+    Interceptor,
+    WorkflowInboundInterceptor,
+    WorkflowInterceptorClassInput,
+)
+
+from application_sdk.activities.common.utils import build_output_path
+from application_sdk.constants import CLEANUP_BASE_PATHS, TEMPORARY_PATH
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+activity.logger = logger
+workflow.logger = logger
+
+
+class CleanupResult(BaseModel):
+    """Result model for cleanup operations.
+
+    Attributes:
+        path_results (Dict[str, bool]): Cleanup results for each path (True=success, False=failure)
+    """
+
+    path_results: Dict[str, bool]
+
+
+@activity.defn
+async def cleanup() -> CleanupResult:
+    """Clean up temporary artifacts and activity state for the current workflow.
+
+    Performs two types of cleanup:
+    1. File cleanup: Removes all contents from configured base paths or default workflow directory
+    2. State cleanup: Clears activity state for the current workflow (includes resource cleanup)
+
+    Uses CLEANUP_BASE_PATHS constant or defaults to workflow-specific artifacts directory.
+
+    Returns:
+        CleanupResult: Structured cleanup results with path results and summary statistics.
+    """
+    path_results: Dict[str, bool] = {}
+    base_paths: List[str] = [os.path.join(TEMPORARY_PATH, build_output_path())]
+
+    # Use configured paths or default to workflow-specific artifacts directory
+    if CLEANUP_BASE_PATHS:
+        base_paths = CLEANUP_BASE_PATHS
+        logger.info(f"Using CLEANUP_BASE_PATHS: {base_paths} for cleanup")
+
+    logger.info(f"Cleaning up all contents from base paths: {base_paths}")
+
+    for base_path in base_paths:
+        try:
+            if os.path.exists(base_path):
+                if os.path.isdir(base_path):
+                    # Remove entire directory and recreate it empty
+                    shutil.rmtree(base_path)
+                    logger.info(f"Cleaned up all contents from: {base_path}")
+                    path_results[base_path] = True
+                else:
+                    logger.warning(f"Path is not a directory: {base_path}")
+                    path_results[base_path] = False
+            else:
+                logger.debug(f"Directory doesn't exist: {base_path}")
+                path_results[base_path] = True
+
+        except Exception as e:
+            logger.error(f"Unexpected error cleaning up {base_path}: {e}")
+            path_results[base_path] = False
+
+    return CleanupResult(
+        path_results=path_results,
+    )
+
+
+class CleanupWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+    """Interceptor for workflow-level app artifacts cleanup.
+
+    This interceptor cleans up the entire app directory structure when the workflow
+    completes or fails, following the pattern: base_path/appname/workflow_id/run_id
+    Supports multiple base paths for comprehensive cleanup.
+    """
+
+    async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
+        """Execute a workflow with app artifacts cleanup.
+
+        Args:
+            input (ExecuteWorkflowInput): The workflow execution input
+
+        Returns:
+            Any: The result of the workflow execution
+
+        Raises:
+            Exception: Re-raises any exceptions from workflow execution
+        """
+        output = None
+        try:
+            output = await super().execute_workflow(input)
+        except Exception:
+            raise
+
+        finally:
+            # Always attempt cleanup regardless of workflow success/failure
+            try:
+                await workflow.execute_activity(
+                    cleanup,
+                    schedule_to_close_timeout=timedelta(minutes=5),
+                    retry_policy=RetryPolicy(
+                        maximum_attempts=3,
+                    ),
+                    summary="This activity is used to cleanup the local artifacts and the activity state after the workflow is completed.",
+                )
+
+                logger.info("Cleanup completed successfully")
+
+            except Exception as e:
+                logger.warning(f"Failed to cleanup artifacts: {e}")
+                # Don't re-raise - cleanup failures shouldn't fail the workflow
+
+        return output
+
+
+class CleanupInterceptor(Interceptor):
+    """Temporal interceptor for automatic app artifacts cleanup.
+
+    This interceptor provides cleanup capabilities for application artifacts
+    across multiple base paths following the pattern: base_path/appname/workflow_id/run_id
+
+    Features:
+    - Automatic cleanup of app-specific artifact directories
+    - Cleanup on workflow completion or failure
+    - Supports multiple cleanup paths via ATLAN_CLEANUP_BASE_PATHS env var
+    - Simple activity-based cleanup logic
+    - Comprehensive error handling and logging
+
+    Example:
+        >>> # Register the interceptor with Temporal worker
+        >>> worker = Worker(
+        ...     client,
+        ...     task_queue="my-task-queue",
+        ...     workflows=[MyWorkflow],
+        ...     activities=[my_activity, cleanup],
+        ...     interceptors=[CleanupInterceptor()]
+        ... )
+
+    Environment Configuration:
+        >>> # Single path (default)
+        >>> ATLAN_CLEANUP_BASE_PATHS="./local/tmp/artifacts/apps"
+
+        >>> # Multiple paths (comma-separated)
+        >>> ATLAN_CLEANUP_BASE_PATHS="./local/tmp/artifacts/apps,/storage/temp/apps,/shared/cleanup/apps"
+    """
+
+    def workflow_interceptor_class(
+        self, input: WorkflowInterceptorClassInput
+    ) -> Optional[Type[WorkflowInboundInterceptor]]:
+        """Get the workflow interceptor class for cleanup.
+
+        Args:
+            input (WorkflowInterceptorClassInput): The interceptor input
+
+        Returns:
+            Optional[Type[WorkflowInboundInterceptor]]: The workflow interceptor class
+        """
+        return CleanupWorkflowInboundInterceptor

application_sdk/interceptors/events.py

@@ -23,6 +23,8 @@ from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.services.eventstore import EventStore
 
 logger = get_logger(__name__)
+activity.logger = logger
+workflow.logger = logger
 
 TEMPORAL_NOT_FOUND_FAILURE = (
     "type.googleapis.com/temporal.api.errordetails.v1.NotFoundFailure"
@@ -41,9 +43,9 @@ async def publish_event(event_data: dict) -> None:
     try:
         event = Event(**event_data)
         await EventStore.publish_event(event)
-        activity.logger.info(f"Published event: {event_data.get('event_name','')}")
+        logger.info(f"Published event: {event_data.get('event_name','')}")
     except Exception as e:
-        activity.logger.error(f"Failed to publish event: {e}")
+        logger.error(f"Failed to publish event: {e}")
         raise
 
 
@@ -123,7 +125,7 @@ class EventWorkflowInboundInterceptor(WorkflowInboundInterceptor):
                 retry_policy=RetryPolicy(maximum_attempts=3),
             )
         except Exception as e:
-            workflow.logger.warning(f"Failed to publish workflow start event: {e}")
+            logger.warning(f"Failed to publish workflow start event: {e}")
             # Don't fail the workflow if event publishing fails
 
         output = None
@@ -152,9 +154,7 @@ class EventWorkflowInboundInterceptor(WorkflowInboundInterceptor):
                     retry_policy=RetryPolicy(maximum_attempts=3),
                 )
             except Exception as publish_error:
-                workflow.logger.warning(
-                    f"Failed to publish workflow end event: {publish_error}"
-                )
+                logger.warning(f"Failed to publish workflow end event: {publish_error}")
 
         return output
 

application_sdk/observability/decorators/observability_decorator.py

@@ -4,7 +4,9 @@ import time
 import uuid
 from typing import Any, Callable, TypeVar, cast
 
-from application_sdk.observability.metrics_adaptor import MetricType
+from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
+from application_sdk.observability.traces_adaptor import get_traces
 
 T = TypeVar("T")
 
@@ -136,9 +138,9 @@ def _record_error_observability(
 
 
 def observability(
-    logger: Any,
-    metrics: Any,
-    traces: Any,
+    logger: Any = None,
+    metrics: Any = None,
+    traces: Any = None,
 ) -> Callable[[Callable[..., T]], Callable[..., T]]:
     """Decorator for adding observability to functions.
 
@@ -146,16 +148,23 @@ def observability(
     It handles both synchronous and asynchronous functions.
 
     Args:
-        logger: Logger instance for operation logging
-        metrics: Metrics adapter for recording operation metrics
-        traces: Traces adapter for recording operation traces
+        logger: Logger instance for operation logging. If None, auto-initializes using get_logger()
+        metrics: Metrics adapter for recording operation metrics. If None, auto-initializes using get_metrics()
+        traces: Traces adapter for recording operation traces. If None, auto-initializes using get_traces()
 
     Returns:
         Callable: Decorated function with observability
 
     Example:
         ```python
+        # With explicit observability components
         @observability(logger, metrics, traces)
+        async def my_function():
+            # Function implementation
+            pass
+
+        # With auto-initialization (recommended)
+        @observability()
         async def my_function():
             # Function implementation
             pass
@@ -163,6 +172,11 @@ def observability(
     """
 
     def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        # Auto-initialize observability components if not provided
+        actual_logger = logger or get_logger(func.__module__)
+        actual_metrics = metrics or get_metrics()
+        actual_traces = traces or get_traces()
+
         # Get function metadata
         func_name = func.__name__
         func_doc = func.__doc__ or f"Executing {func_name}"
@@ -170,7 +184,7 @@ def observability(
         is_async = inspect.iscoroutinefunction(func)
 
         # Debug logging for function decoration
-        logger.debug(f"Decorating function {func_name} (async={is_async})")
+        actual_logger.debug(f"Decorating function {func_name} (async={is_async})")
 
         @functools.wraps(func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> T:
@@ -181,16 +195,16 @@ def observability(
 
             try:
                 # Log start of operation
-                logger.debug(f"Starting async function {func_name}")
+                actual_logger.debug(f"Starting async function {func_name}")
 
                 # Execute the function
                 result = await func(*args, **kwargs)
 
                 # Record success observability
                 _record_success_observability(
-                    logger,
-                    metrics,
-                    traces,
+                    actual_logger,
+                    actual_metrics,
+                    actual_traces,
                     func_name,
                     func_doc,
                     func_module,
@@ -204,9 +218,9 @@ def observability(
             except Exception as e:
                 # Record error observability
                 _record_error_observability(
-                    logger,
-                    metrics,
-                    traces,
+                    actual_logger,
+                    actual_metrics,
+                    actual_traces,
                     func_name,
                     func_doc,
                     func_module,
@@ -226,16 +240,16 @@ def observability(
 
             try:
                 # Log start of operation
-                logger.debug(f"Starting sync function {func_name}")
+                actual_logger.debug(f"Starting sync function {func_name}")
 
                 # Execute the function
                 result = func(*args, **kwargs)
 
                 # Record success observability
                 _record_success_observability(
-                    logger,
-                    metrics,
-                    traces,
+                    actual_logger,
+                    actual_metrics,
+                    actual_traces,
                     func_name,
                     func_doc,
                     func_module,
@@ -249,9 +263,9 @@ def observability(
             except Exception as e:
                 # Record error observability
                 _record_error_observability(
-                    logger,
-                    metrics,
-                    traces,
+                    actual_logger,
+                    actual_metrics,
+                    actual_traces,
                     func_name,
                     func_doc,
                     func_module,

application_sdk/outputs/.cursor/BUGBOT.md

@@ -0,0 +1,295 @@
+# Output Code Review Guidelines - Data Output Processing
+
+## Context-Specific Patterns
+
+This directory contains output processing implementations for various data formats (JSON, Parquet, Iceberg). Output processors must handle data uploads efficiently while maintaining data integrity and correct destination paths.
+
+### Phase 1: Critical Output Safety Issues
+
+**Object Store Path Management:**
+
+- **Correct destination paths**: Upload paths must respect user-configured output prefixes
+- **Path construction accuracy**: Object store keys must be calculated correctly, not hardcoded
+- **User prefix preservation**: Respect user-provided output directories and naming conventions
+- **Path validation**: Ensure upload paths don't conflict with existing data
+
+**Data Integrity and Security:**
+
+- All output data must be validated before upload
+- File permissions and access controls must be properly set
+- Data serialization must be consistent and recoverable
+- Prevent overwriting critical data without confirmation
+- Maintain data lineage information in output metadata
+
+```python
+# ✅ DO: Proper object store upload path handling
+class JsonOutput:
+    async def upload_to_object_store(
+        self,
+        data: List[dict],
+        output_prefix: str,  # User-provided output location
+        filename: str
+    ) -> dict:
+        """Upload data with correct path handling."""
+
+        # Construct full object store path respecting user's output prefix
+        object_store_key = os.path.join(output_prefix, filename)
+
+        # Serialize data
+        json_data = orjson.dumps(data, option=orjson.OPT_APPEND_NEWLINE)
+
+        # Upload to correct location
+        result = await self.object_store.upload_file(
+            data=json_data,
+            destination=object_store_key  # Respect user's intended location
+        )
+
+        return result
+
+# ❌ REJECT: Incorrect path handling
+class BadJsonOutput:
+    async def upload_to_object_store(self, data: List[dict], filename: str):
+        # Wrong: hardcoded or derived path, ignoring user configuration
+        object_store_key = get_object_store_prefix(f"/tmp/{filename}")  # Ignores output_prefix!
+
+        result = await self.object_store.upload_file(
+            data=orjson.dumps(data),
+            destination=object_store_key  # Wrong destination!
+        )
+        return result
+```
+
+### Phase 2: Output Architecture Patterns
+
+**Performance Optimization Requirements:**
+
+- **Parallelization opportunities**: Flag sequential upload operations that could be parallelized
+- **Batch processing**: Group related uploads to reduce overhead
+- **Streaming uploads**: Use streaming for large datasets instead of loading into memory
+- **Connection optimization**: Reuse object store connections across operations
+
+**Resource Management:**
+
+- Use proper connection pooling for object store operations
+- Implement timeout handling for upload operations
+- Clean up temporary files after upload
+- Handle partial upload failures gracefully
+- Monitor memory usage during large data serialization
+
+```python
+# ✅ DO: Parallel upload processing
+async def upload_multiple_datasets_parallel(
+    self,
+    datasets: List[Tuple[List[dict], str]],  # (data, filename) pairs
+    output_prefix: str
+) -> List[dict]:
+    """Upload multiple datasets in parallel for better performance."""
+
+    async def upload_single_dataset(data: List[dict], filename: str) -> dict:
+        """Upload a single dataset with error handling."""
+        try:
+            object_store_key = os.path.join(output_prefix, filename)
+            serialized_data = orjson.dumps(data, option=orjson.OPT_APPEND_NEWLINE)
+
+            return await self.object_store.upload_file(
+                data=serialized_data,
+                destination=object_store_key
+            )
+        except Exception as e:
+            logger.error(f"Failed to upload {filename}: {e}")
+            raise
+
+    # Parallel processing with controlled concurrency
+    semaphore = asyncio.Semaphore(5)  # Limit concurrent uploads
+
+    async def upload_with_semaphore(data: List[dict], filename: str) -> dict:
+        async with semaphore:
+            return await upload_single_dataset(data, filename)
+
+    tasks = [upload_with_semaphore(data, filename) for data, filename in datasets]
+    return await asyncio.gather(*tasks)
+
+# ❌ REJECT: Sequential upload processing
+async def upload_multiple_datasets_sequential(
+    self,
+    datasets: List[Tuple[List[dict], str]],
+    output_prefix: str
+) -> List[dict]:
+    """Sequential uploads - should be flagged for parallelization."""
+    results = []
+    for data, filename in datasets:  # FLAG: Could be parallelized
+        object_store_key = os.path.join(output_prefix, filename)
+        result = await self.object_store.upload_file(data, object_store_key)
+        results.append(result)
+    return results
+```
+
+### Phase 3: Output Testing Requirements
+
+**Data Output Testing:**
+
+- Test with various data formats and sizes
+- Test serialization and deserialization consistency
+- Test partial upload scenarios and recovery
+- Mock object store operations in unit tests
+- Include integration tests with real object store
+- Test data corruption detection and prevention
+
+**Performance Testing:**
+
+- Include tests for large dataset uploads
+- Test memory usage during serialization
+- Test concurrent upload operations
+- Verify timeout handling works correctly
+- Test connection pool behavior under load
+
+### Phase 4: Performance and Scalability
+
+**Data Upload Efficiency:**
+
+- Use streaming uploads for large datasets
+- Implement proper chunking for oversized data
+- Use compression for large text-based outputs
+- Monitor upload progress and provide feedback
+- Optimize serialization performance (use orjson over json)
+
+**Object Store Optimization:**
+
+- Use connection pooling for object store clients
+- Implement proper retry logic for upload failures
+- Use parallel uploads where appropriate
+- Monitor upload metrics and error rates
+- Handle bandwidth limitations gracefully
+
+### Phase 5: Output Maintainability
+
+**Error Handling and Recovery:**
+
+- Implement comprehensive error handling for all upload operations
+- Provide meaningful error messages with upload context
+- Handle partial upload failures gracefully
+- Implement proper retry logic for transient failures
+- Log all upload operations with destination information
+
+**Configuration Management:**
+
+- Externalize all output-related configuration
+- Support different output destinations and formats
+- Validate output configuration before processing
+- Document all supported output parameters
+- Handle environment-specific output requirements
+
+---
+
+## Output-Specific Anti-Patterns
+
+**Always Reject:**
+
+- **Path derivation errors**: Deriving object store paths from local temporary paths
+- **Sequential uploads**: Uploading multiple files sequentially when parallel uploads are possible
+- **Memory inefficiency**: Loading entire datasets into memory for serialization
+- **Missing upload verification**: Not verifying successful uploads
+- **Poor error recovery**: Not handling partial upload failures gracefully
+- **Resource leaks**: Not cleaning up temporary files or connections
+
+**Object Store Upload Anti-Patterns:**
+
+```python
+# ❌ REJECT: Incorrect upload path handling
+class BadOutputProcessor:
+    async def upload_results(self, results: List[dict]):
+        # Wrong: derives upload path from temporary local path
+        local_temp_file = "/tmp/results.json"
+        upload_key = get_object_store_prefix(local_temp_file)  # Incorrect!
+
+        await self.object_store.upload_file(results, upload_key)
+
+# ✅ REQUIRE: Correct upload path handling
+class GoodOutputProcessor:
+    async def upload_results(
+        self,
+        results: List[dict],
+        output_prefix: str,  # User-specified destination
+        filename: str = "results.json"
+    ):
+        # Use actual user-configured output location
+        upload_key = os.path.join(output_prefix, filename)
+
+        await self.object_store.upload_file(
+            data=orjson.dumps(results),
+            destination=upload_key  # Correct destination
+        )
+```
+
+**Performance Anti-Patterns:**
+
+```python
+# ❌ REJECT: Sequential upload processing
+async def upload_multiple_files_sequential(file_data_pairs: List[Tuple]):
+    results = []
+    for data, filename in file_data_pairs:  # Should be parallelized
+        result = await upload_single_file(data, filename)
+        results.append(result)
+    return results
+
+# ✅ REQUIRE: Parallel upload processing with proper error handling
+async def upload_multiple_files_parallel(
+    file_data_pairs: List[Tuple],
+    max_concurrency: int = 5
+) -> List[dict]:
+    semaphore = asyncio.Semaphore(max_concurrency)
+
+    async def upload_with_semaphore(data, filename):
+        async with semaphore:
+            try:
+                return await upload_single_file(data, filename)
+            except Exception as e:
+                logger.error(f"Upload failed for {filename}: {e}")
+                return {"filename": filename, "status": "failed", "error": str(e)}
+
+    tasks = [upload_with_semaphore(data, filename) for data, filename in file_data_pairs]
+    return await asyncio.gather(*tasks)
+```
+
+**Memory Management Anti-Patterns:**
+
+```python
+# ❌ REJECT: Loading entire dataset for serialization
+async def bad_large_dataset_upload(large_dataset: List[dict]):
+    # Loads entire dataset into memory
+    json_data = orjson.dumps(large_dataset)  # Could exceed memory limits
+    await upload_data(json_data)
+
+# ✅ REQUIRE: Streaming serialization for large datasets
+async def good_large_dataset_upload(large_dataset: List[dict], chunk_size: int = 1000):
+    """Stream large datasets to avoid memory issues."""
+
+    async def serialize_chunk(chunk: List[dict]) -> bytes:
+        return orjson.dumps(chunk, option=orjson.OPT_APPEND_NEWLINE)
+
+    # Process in chunks to manage memory
+    for i in range(0, len(large_dataset), chunk_size):
+        chunk = large_dataset[i:i + chunk_size]
+        serialized_chunk = await serialize_chunk(chunk)
+
+        await upload_chunk(
+            data=serialized_chunk,
+            chunk_index=i // chunk_size
+        )
+```
+
+## Educational Context for Output Reviews
+
+When reviewing output code, emphasize:
+
+1. **Data Integrity Impact**: "Incorrect upload path handling can cause data to be stored in wrong locations, making it inaccessible to downstream processes. This breaks data pipelines and can cause data loss."
+
+2. **Performance Impact**: "Sequential uploads create unnecessary bottlenecks. For enterprise datasets with multiple output files, parallelization can significantly reduce processing time and improve user experience."
+
+3. **Resource Impact**: "Poor memory management during serialization can cause out-of-memory errors with large datasets. Streaming and chunking are essential for enterprise-scale data output."
+
+4. **User Experience Impact**: "Output path errors are often discovered late in processing, causing wasted computation and frustrating delays. Proper validation and clear error messages improve reliability."
+
+5. **Scalability Impact**: "Output patterns that work for small datasets can fail at enterprise scale. Always design output processes to handle the largest expected dataset sizes efficiently."
+
+6. **Data Pipeline Impact**: "Output processing is the final step in data pipelines. Failures here can invalidate all upstream processing work. Robust error handling and verification are critical for pipeline reliability."

application_sdk/outputs/iceberg.py

@@ -29,6 +29,7 @@ class IcebergOutput(Output):
         mode: str = "append",
         total_record_count: int = 0,
         chunk_count: int = 0,
+        retain_local_copy: bool = False,
     ):
         """Initialize the Iceberg output class.
 
@@ -39,6 +40,8 @@ class IcebergOutput(Output):
             mode (str, optional): Write mode for the iceberg table. Defaults to "append".
             total_record_count (int, optional): Total record count written to the iceberg table. Defaults to 0.
             chunk_count (int, optional): Number of chunks written to the iceberg table. Defaults to 0.
+            retain_local_copy (bool, optional): Whether to retain the local copy of the files.
+                Defaults to False.
         """
         self.total_record_count = total_record_count
         self.chunk_count = chunk_count
@@ -47,6 +50,7 @@ class IcebergOutput(Output):
         self.iceberg_table = iceberg_table
         self.mode = mode
         self.metrics = get_metrics()
+        self.retain_local_copy = retain_local_copy
 
     async def write_dataframe(self, dataframe: "pd.DataFrame"):
         """