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

application_sdk/common/utils.py

@@ -17,8 +17,12 @@ from typing import (
     Union,
 )
 
+from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.common.error_codes import CommonError
+from application_sdk.constants import TEMPORARY_PATH
+from application_sdk.inputs.sql_query import SQLQueryInput
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 
@@ -106,10 +110,42 @@ def extract_database_names_from_regex_common(
         return empty_default
 
 
+def transform_posix_regex(regex_pattern: str) -> str:
+    r"""
+    Transform regex pattern for POSIX compatibility.
+
+    Rules:
+    1. Add ^ before each database name before \.
+    2. Add an additional . between \. and * if * follows \.
+
+    Example: 'dev\.public$|dev\.atlan_test_schema$|wide_world_importers\.*'
+    Becomes: '^dev\.public$|^dev\.atlan_test_schema$|^wide_world_importers\..*'
+    """
+    if not regex_pattern:
+        return regex_pattern
+
+    # Split by | to handle each pattern separately
+    patterns = regex_pattern.split("|")
+    transformed_patterns = []
+
+    for pattern in patterns:
+        # Add ^ at the beginning if it's not already there
+        if not pattern.startswith("^"):
+            pattern = "^" + pattern
+
+            # Add additional . between \. and * if * follows \.
+            pattern = re.sub(r"\\\.\*", r"\..*", pattern)
+
+        transformed_patterns.append(pattern)
+
+    return "|".join(transformed_patterns)
+
+
 def prepare_query(
     query: Optional[str],
     workflow_args: Dict[str, Any],
     temp_table_regex_sql: Optional[str] = "",
+    use_posix_regex: Optional[bool] = False,
 ) -> Optional[str]:
     """
     Prepares a SQL query by applying include and exclude filters, and optional
@@ -158,6 +194,14 @@ def prepare_query(
             include_filter, exclude_filter
         )
 
+        if use_posix_regex:
+            normalized_include_regex_posix = transform_posix_regex(
+                normalized_include_regex
+            )
+            normalized_exclude_regex_posix = transform_posix_regex(
+                normalized_exclude_regex
+            )
+
         # Extract database names from the normalized regex patterns
         include_databases = extract_database_names_from_regex_common(
             normalized_regex=normalized_include_regex,
@@ -176,15 +220,26 @@ def prepare_query(
         )
         exclude_views = workflow_args.get("metadata", {}).get("exclude_views", False)
 
-        return query.format(
-            include_databases=include_databases,
-            exclude_databases=exclude_databases,
-            normalized_include_regex=normalized_include_regex,
-            normalized_exclude_regex=normalized_exclude_regex,
-            temp_table_regex_sql=temp_table_regex_sql,
-            exclude_empty_tables=exclude_empty_tables,
-            exclude_views=exclude_views,
-        )
+        if use_posix_regex:
+            return query.format(
+                include_databases=include_databases,
+                exclude_databases=exclude_databases,
+                normalized_include_regex=normalized_include_regex_posix,
+                normalized_exclude_regex=normalized_exclude_regex_posix,
+                temp_table_regex_sql=temp_table_regex_sql,
+                exclude_empty_tables=exclude_empty_tables,
+                exclude_views=exclude_views,
+            )
+        else:
+            return query.format(
+                include_databases=include_databases,
+                exclude_databases=exclude_databases,
+                normalized_include_regex=normalized_include_regex,
+                normalized_exclude_regex=normalized_exclude_regex,
+                temp_table_regex_sql=temp_table_regex_sql,
+                exclude_empty_tables=exclude_empty_tables,
+                exclude_views=exclude_views,
+            )
     except CommonError as e:
         # Extract the original error message from the CommonError
         error_message = str(e).split(": ", 1)[-1] if ": " in str(e) else str(e)
@@ -195,6 +250,47 @@ def prepare_query(
         return None
 
 
+async def get_database_names(
+    sql_client, workflow_args, fetch_database_sql
+) -> Optional[List[str]]:
+    """
+    Get the database names from the workflow args if include-filter is present
+    Args:
+        workflow_args: The workflow args
+    Returns:
+        List[str]: The database names
+    """
+    database_names = parse_filter_input(
+        workflow_args.get("metadata", {}).get("include-filter", {})
+    )
+
+    database_names = [
+        re.sub(r"^[^\w]+|[^\w]+$", "", database_name)
+        for database_name in database_names
+    ]
+    if not database_names:
+        # if database_names are not provided in the include-filter, we'll run the query to get all the database names
+        # because by default for an empty include-filter, we fetch details corresponding to all the databases.
+        temp_table_regex_sql = workflow_args.get("metadata", {}).get(
+            "temp-table-regex", ""
+        )
+        prepared_query = prepare_query(
+            query=fetch_database_sql,
+            workflow_args=workflow_args,
+            temp_table_regex_sql=temp_table_regex_sql,
+            use_posix_regex=True,
+        )
+        # We'll run the query to get all the database names
+        database_sql_input = SQLQueryInput(
+            engine=sql_client.engine,
+            query=prepared_query,  # type: ignore
+            chunk_size=None,
+        )
+        database_dataframe = await database_sql_input.get_dataframe()
+        database_names = list(database_dataframe["database_name"])
+    return database_names
+
+
 def parse_filter_input(
     filter_input: Union[str, Dict[str, Any], None],
 ) -> Dict[str, Any]:
@@ -416,6 +512,46 @@ def parse_credentials_extra(credentials: Dict[str, Any]) -> Dict[str, Any]:
     return extra  # We know it's a Dict[str, Any] due to the Union type and str check
 
 
+def has_custom_control_config(workflow_args: Dict[str, Any]) -> bool:
+    """
+    Check if custom control configuration is present in workflow arguments.
+
+    Args:
+        workflow_args: The workflow arguments
+
+    Returns:
+        bool: True if custom control configuration is present, False otherwise
+    """
+    return (
+        workflow_args.get("control-config-strategy") == "custom"
+        and workflow_args.get("control-config") is not None
+    )
+
+
+async def get_file_names(output_path: str, typename: str) -> List[str]:
+    """
+    Get file names for a specific asset type from the transformed directory.
+
+    Args:
+        output_path (str): The base output path
+        typename (str): The asset type (e.g., 'table', 'schema', 'column')
+
+    Returns:
+        List[str]: List of relative file paths for the asset type
+    """
+
+    source = get_object_store_prefix(os.path.join(output_path, typename))
+    await ObjectStore.download_prefix(source, TEMPORARY_PATH)
+
+    file_pattern = os.path.join(output_path, typename, "*.json")
+    file_names = glob.glob(file_pattern)
+    file_name_list = [
+        "/".join(file_name.rsplit("/", 2)[-2:]) for file_name in file_names
+    ]
+
+    return file_name_list
+
+
 def run_sync(func):
     """Run a function in a thread pool executor.
 

application_sdk/constants.py

@@ -59,6 +59,14 @@ WORKFLOW_OUTPUT_PATH_TEMPLATE = (
 # Temporary Path (used to store intermediate files)
 TEMPORARY_PATH = os.getenv("ATLAN_TEMPORARY_PATH", "./local/tmp/")
 
+# Cleanup Paths (custom paths for cleanup operations, supports multiple paths separated by comma)
+# If empty, cleanup activities will default to workflow-specific paths at runtime
+CLEANUP_BASE_PATHS = [
+    path.strip()
+    for path in os.getenv("ATLAN_CLEANUP_BASE_PATHS", "").split(",")
+    if path.strip()
+]
+
 # State Store Constants
 #: Path template for state store files (example: objectstore://bucket/persistent-artifacts/apps/{application_name}/{state_type}/{id}/config.json)
 STATE_STORE_PATH_TEMPLATE = (

application_sdk/decorators/.cursor/BUGBOT.md

@@ -0,0 +1,279 @@
+# Decorator Code Review Guidelines - Centralized Function Decorators
+
+## Context-Specific Patterns
+
+This directory contains all decorator implementations for the Application SDK. Decorators must be centralized here to avoid scattered functionality and ensure consistent patterns.
+
+### Phase 1: Critical Decorator Safety Issues
+
+**Decorator Centralization:**
+
+- **ALL decorators must be in this directory**: No decorators should exist in other modules (lock/, observability/, etc.)
+- **Consolidate scattered decorators**: If decorators are found elsewhere, they must be moved here
+- **Single responsibility per file**: Each decorator type should have its own file (locks.py, observability_decorator.py)
+- **Proper imports**: Other modules should import decorators from here, not define their own
+
+**Type Safety and Function Signatures:**
+
+- All decorators must preserve function signatures and type hints
+- Use `functools.wraps` to maintain function metadata
+- Generic decorators must use proper type annotations
+- Return types must match the original function's return type
+
+```python
+# ✅ DO: Proper decorator type safety
+from typing import Callable, Any, TypeVar, ParamSpec
+from functools import wraps
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+def my_decorator(func: Callable[P, T]) -> Callable[P, T]:
+    @wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        # Decorator logic
+        return func(*args, **kwargs)
+    return wrapper
+
+# ❌ NEVER: Poor type annotations
+def bad_decorator(func):  # No type hints
+    def wrapper(*args, **kwargs):  # No type preservation
+        return func(*args, **kwargs)
+    return wrapper  # Missing @wraps
+```
+
+### Phase 2: Decorator Architecture Patterns
+
+**Proper Decorator Structure:**
+
+- **Parameterized decorators**: Support both `@decorator` and `@decorator(param=value)` usage patterns
+- **Error handling**: Decorators must not swallow exceptions unless explicitly designed to do so
+- **Resource cleanup**: Decorators that acquire resources must ensure cleanup in finally blocks
+- **Context preservation**: Maintain original function context and metadata
+
+**Configuration Management:**
+
+- **Centralized constants**: All decorator configuration should use constants from this directory
+- **Shared configuration**: Related decorators should share configuration patterns
+- **Environment awareness**: Decorators should work in both development and production environments
+
+```python
+# ✅ DO: Proper decorator configuration
+from application_sdk.constants import DEFAULT_LOCK_TTL, DEFAULT_MAX_LOCKS
+
+# Shared configuration for lock decorators
+LOCK_CONFIG_KEY = "distributed_lock_config"  # Centralized key
+
+def distributed_lock(
+    lock_name: Optional[str] = None,
+    max_locks: int = DEFAULT_MAX_LOCKS,
+    ttl_seconds: int = DEFAULT_LOCK_TTL
+):
+    """Distributed lock decorator with proper defaults and configuration."""
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        @wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            # Use centralized configuration
+            actual_lock_name = lock_name or f"{func.__module__}.{func.__name__}"
+
+            # Store config in activity context using shared key
+            activity_info = activity.info()
+            activity_info.memo[LOCK_CONFIG_KEY] = {
+                "lock_name": actual_lock_name,
+                "max_locks": max_locks,
+                "ttl_seconds": ttl_seconds
+            }
+
+            return await func(*args, **kwargs)
+        return wrapper
+    return decorator
+
+# ❌ REJECT: Scattered constants and configuration
+def bad_lock_decorator(max_locks=10):  # Hardcoded default
+    LOCK_KEY = "my_lock_key"  # Should be centralized
+    # Configuration scattered across files
+```
+
+### Phase 3: Decorator Testing Requirements
+
+**Comprehensive Decorator Testing:**
+
+- **Function preservation**: Test that decorators preserve original function behavior
+- **Type safety**: Verify type hints are maintained after decoration
+- **Error propagation**: Ensure exceptions are properly handled and propagated
+- **Resource cleanup**: Test cleanup behavior in both success and failure cases
+- **Configuration validation**: Test all configuration parameters and edge cases
+
+```python
+# ✅ DO: Comprehensive decorator testing
+@pytest.mark.asyncio
+class TestDistributedLockDecorator:
+    """Test suite for distributed lock decorator."""
+
+    async def test_function_signature_preservation(self):
+        """Test that decorator preserves function signature and types."""
+
+        @distributed_lock("test_lock")
+        async def test_function(param1: str, param2: int = 10) -> dict:
+            """Test function docstring."""
+            return {"param1": param1, "param2": param2}
+
+        # Verify signature preservation
+        assert test_function.__name__ == "test_function"
+        assert test_function.__doc__ == "Test function docstring."
+
+        # Verify function still works
+        result = await test_function("test", 20)
+        assert result == {"param1": "test", "param2": 20}
+
+    async def test_error_propagation(self):
+        """Test that decorator properly propagates exceptions."""
+
+        @distributed_lock("error_lock")
+        async def failing_function():
+            raise ValueError("Test error")
+
+        # Verify exception is propagated, not swallowed
+        with pytest.raises(ValueError, match="Test error"):
+            await failing_function()
+
+    async def test_resource_cleanup_on_failure(self, mock_lock_manager):
+        """Test that resources are cleaned up even when function fails."""
+
+        @distributed_lock("cleanup_test")
+        async def failing_function():
+            raise RuntimeError("Simulated failure")
+
+        mock_lock_manager.acquire_lock.return_value.__aenter__ = AsyncMock()
+        mock_lock_manager.acquire_lock.return_value.__aexit__ = AsyncMock()
+
+        with pytest.raises(RuntimeError):
+            await failing_function()
+
+        # Verify cleanup was called
+        mock_lock_manager.acquire_lock.return_value.__aexit__.assert_called_once()
+```
+
+### Phase 4: Performance and Integration
+
+**Decorator Performance:**
+
+- **Minimal overhead**: Decorators should add minimal performance overhead
+- **Async compatibility**: All decorators must work correctly with async functions
+- **Context manager efficiency**: Use efficient context managers for resource management
+- **Caching**: Cache expensive decorator setup operations where appropriate
+
+**Integration Patterns:**
+
+- **Temporal integration**: Decorators must work correctly with Temporal activities and workflows
+- **Observability integration**: Integrate with logging, metrics, and tracing systems
+- **Error handling integration**: Work correctly with the SDK's error handling patterns
+
+### Phase 5: Decorator Maintainability
+
+**Code Organization:**
+
+- **One decorator type per file**: Keep related decorators together (all lock decorators in locks.py)
+- **Clear naming**: Decorator files should clearly indicate their purpose
+- **Consistent patterns**: All decorators should follow the same structural patterns
+- **Documentation**: Each decorator must have comprehensive docstrings with usage examples
+
+**Backwards Compatibility:**
+
+- **API stability**: Decorator APIs should be stable across versions
+- **Graceful deprecation**: Deprecated decorators should include migration guidance
+- **Version compatibility**: Support existing usage patterns when adding new features
+
+---
+
+## Decorator-Specific Anti-Patterns
+
+**Always Reject:**
+
+- **Scattered decorators**: Decorators defined outside this directory
+- **Missing type safety**: Decorators without proper type annotations
+- **Resource leaks**: Decorators that don't clean up resources properly
+- **Exception swallowing**: Decorators that hide exceptions unintentionally
+- **Poor configuration**: Hardcoded values that should be configurable
+- **No function preservation**: Decorators that don't preserve original function metadata
+
+**Centralization Anti-Patterns:**
+
+```python
+# ❌ REJECT: Decorators in wrong locations
+# Found in application_sdk/lock/__init__.py
+def needs_lock(max_locks=10):
+    """Should be in decorators/locks.py instead"""
+
+# Found in application_sdk/observability/some_module.py
+def trace_activity(func):
+    """Should be in decorators/observability_decorator.py"""
+
+# ✅ REQUIRE: Centralized decorators
+# In application_sdk/decorators/locks.py
+def needs_lock(max_locks: int = DEFAULT_MAX_LOCKS):
+    """Properly located distributed lock decorator"""
+
+# In application_sdk/decorators/observability_decorator.py
+def observability(logger=None, metrics=None, traces=None):
+    """Properly located observability decorator"""
+```
+
+**Type Safety Anti-Patterns:**
+
+```python
+# ❌ REJECT: Poor type safety
+def bad_decorator(func):  # No type annotations
+    def wrapper(*args, **kwargs):  # No parameter specifications
+        return func(*args, **kwargs)
+    return wrapper  # Missing @wraps, no return type
+
+# ✅ REQUIRE: Proper type safety
+from typing import Callable, TypeVar, ParamSpec
+from functools import wraps
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+def good_decorator(func: Callable[P, T]) -> Callable[P, T]:
+    @wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        return func(*args, **kwargs)
+    return wrapper
+```
+
+**Configuration Anti-Patterns:**
+
+```python
+# ❌ REJECT: Scattered configuration
+# Different files using different keys for same concept
+LOCK_KEY_1 = "lock_config"  # In locks.py
+LOCK_KEY_2 = "distributed_lock"  # In interceptors.py
+DEFAULT_TTL = 300  # Hardcoded in decorator
+
+# ✅ REQUIRE: Centralized configuration
+# In application_sdk/constants.py
+DISTRIBUTED_LOCK_CONFIG_KEY = "distributed_lock_config"
+DEFAULT_LOCK_TTL = 300
+DEFAULT_MAX_LOCKS = 10
+
+# In decorators using shared constants
+from application_sdk.constants import DISTRIBUTED_LOCK_CONFIG_KEY, DEFAULT_LOCK_TTL
+```
+
+## Educational Context for Decorator Reviews
+
+When reviewing decorator code, emphasize:
+
+1. **Centralization Impact**: "Scattered decorators create maintenance nightmares. When the same decorator logic appears in multiple places, bugs get fixed in some places but not others. Centralization ensures consistency and reduces maintenance burden."
+
+2. **Type Safety Impact**: "Decorators that don't preserve type information break IDE support, static analysis, and developer productivity. Proper type annotations are essential for maintaining code quality in large codebases."
+
+3. **Resource Management Impact**: "Decorators often manage resources (locks, connections, contexts). Poor resource management in decorators can cause system-wide issues because they're used across many functions."
+
+4. **Function Preservation Impact**: "Decorators that don't preserve original function metadata break debugging, introspection, and documentation tools. Using @functools.wraps is not optional."
+
+5. **Testing Impact**: "Decorators are cross-cutting concerns that affect many functions. Bugs in decorators have amplified impact, making thorough testing especially critical."
+
+6. **Performance Impact**: "Decorators add overhead to every function call they wrap. Inefficient decorators can degrade system performance across the entire application."

application_sdk/handlers/__init__.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Dict
 
 
 class HandlerInterface(ABC):
@@ -37,3 +37,10 @@ class HandlerInterface(ABC):
         To be implemented by the subclass
         """
         raise NotImplementedError("fetch_metadata method not implemented")
+
+    @staticmethod
+    async def get_configmap(config_map_id: str) -> Dict[str, Any]:
+        """
+        Static method to get the configmap
+        """
+        return {}

application_sdk/handlers/sql.py

@@ -56,9 +56,13 @@ class BaseSQLHandler(HandlerInterface):
     schema_alias_key: str = SQLConstants.SCHEMA_ALIAS_KEY.value
     database_result_key: str = SQLConstants.DATABASE_RESULT_KEY.value
     schema_result_key: str = SQLConstants.SCHEMA_RESULT_KEY.value
+    multidb: bool = False
 
-    def __init__(self, sql_client: BaseSQLClient | None = None):
+    def __init__(
+        self, sql_client: BaseSQLClient | None = None, multidb: Optional[bool] = False
+    ):
         self.sql_client = sql_client
+        self.multidb = multidb
 
     async def load(self, credentials: Dict[str, Any]) -> None:
         """
@@ -294,35 +298,26 @@ class BaseSQLHandler(HandlerInterface):
                         return False, f"{db}.{sch} schema"
         return True, ""
 
-    async def tables_check(
-        self,
-        payload: Dict[str, Any],
-    ) -> Dict[str, Any]:
+    async def tables_check(self, payload: Dict[str, Any]) -> Dict[str, Any]:
         """
         Method to check the count of tables
         """
         logger.info("Starting tables check")
-        query = prepare_query(
-            query=self.tables_check_sql,
-            workflow_args=payload,
-            temp_table_regex_sql=self.extract_temp_table_regex_table_sql,
-        )
-        if not query:
-            raise ValueError("tables_check_sql is not defined")
-        sql_input = SQLQueryInput(
-            engine=self.sql_client.engine, query=query, chunk_size=None
-        )
-        sql_input = await sql_input.get_dataframe()
-        try:
-            result = 0
-            for row in sql_input.to_dict(orient="records"):
-                result += row["count"]
+
+        def _sum_counts_from_records(records_iter) -> int:
+            total = 0
+            for row in records_iter:
+                total += row["count"]
+            return total
+
+        def _build_success(total: int) -> Dict[str, Any]:
             return {
                 "success": True,
-                "successMessage": f"Tables check successful. Table count: {result}",
+                "successMessage": f"Tables check successful. Table count: {total}",
                 "failureMessage": "",
             }
-        except Exception as exc:
+
+        def _build_failure(exc: Exception) -> Dict[str, Any]:
             logger.error("Error during tables check", exc_info=True)
             return {
                 "success": False,
@@ -331,6 +326,52 @@ class BaseSQLHandler(HandlerInterface):
                 "error": str(exc),
             }
 
+        if self.multidb:
+            try:
+                from application_sdk.activities.metadata_extraction.sql import (
+                    BaseSQLMetadataExtractionActivities,
+                )
+
+                # Use the base query executor in multidb mode to get concatenated df
+                activities = BaseSQLMetadataExtractionActivities()
+                activities.multidb = True
+                concatenated_df = await activities.query_executor(
+                    sql_engine=self.sql_client.engine if self.sql_client else None,
+                    sql_query=self.tables_check_sql,
+                    workflow_args=payload,
+                    output_suffix="raw/table",
+                    typename="table",
+                    write_to_file=False,
+                    concatenate=True,
+                    return_dataframe=True,
+                    sql_client=self.sql_client,
+                )
+
+                if concatenated_df is None:
+                    return _build_success(0)
+
+                total = int(concatenated_df["count"].sum())  # type: ignore[index]
+                return _build_success(total)
+            except Exception as exc:
+                return _build_failure(exc)
+        else:
+            query = prepare_query(
+                query=self.tables_check_sql,
+                workflow_args=payload,
+                temp_table_regex_sql=self.extract_temp_table_regex_table_sql,
+            )
+            if not query:
+                raise ValueError("tables_check_sql is not defined")
+            sql_input = SQLQueryInput(
+                engine=self.sql_client.engine, query=query, chunk_size=None
+            )
+            sql_input = await sql_input.get_dataframe()
+            try:
+                total = _sum_counts_from_records(sql_input.to_dict(orient="records"))
+                return _build_success(total)
+            except Exception as exc:
+                return _build_failure(exc)
+
     async def check_client_version(self) -> Dict[str, Any]:
         """
         Check if the client version meets the minimum required version.