atlan-application-sdk 1.1.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

application_sdk/io/utils.py

@@ -0,0 +1,307 @@
+import glob
+import os
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, 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.constants import TEMPORARY_PATH
+from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
+
+JSON_FILE_EXTENSION = ".json"
+PARQUET_FILE_EXTENSION = ".parquet"
+
+if TYPE_CHECKING:
+    import daft
+    import pandas as pd
+
+
+logger = get_logger(__name__)
+
+
+def find_local_files_by_extension(
+    path: str,
+    extension: str,
+    file_names: Optional[List[str]] = None,
+) -> List[str]:
+    """Find local files at the specified local path, optionally filtering by file names.
+
+    Args:
+        path (str): Local path to search in (file or directory)
+        extension (str): File extension to filter by (e.g., '.parquet', '.json')
+        file_names (Optional[List[str]]): List of file names (basenames) to filter by, paths are not supported
+
+    Returns:
+        List[str]: List of matching file paths
+
+    Example:
+        >>> find_local_files_by_extension("/data", ".parquet", ["file1.parquet", "file2.parquet"])
+        ['file1.parquet', 'file2.parquet']
+
+        >>> find_local_files_by_extension("/data/single.json", ".json")
+        ['single.json']
+    """
+    if os.path.isfile(path) and path.endswith(extension):
+        # Single file - return it directly
+        return [path]
+
+    elif os.path.isdir(path):
+        # Directory - find all files in directory
+        all_files = glob.glob(
+            os.path.join(path, "**", f"*{extension}"),
+            recursive=True,
+        )
+
+        # Filter by file names if specified
+        if file_names:
+            file_names_set = set(file_names)  # Convert to set for O(1) lookup
+            return [f for f in all_files if os.path.basename(f) in file_names_set]
+        else:
+            return all_files
+
+    return []
+
+
+async def download_files(
+    path: str, file_extension: str, file_names: Optional[List[str]] = None
+) -> List[str]:
+    """Download files from object store if not available locally.
+
+    Flow:
+    1. Check if files exist locally at self.path
+    2. If not, try to download from object store
+    3. Filter by self.file_names if provided
+    4. Return list of file paths for logging purposes
+
+    Returns:
+        List[str]: List of file paths
+
+    Raises:
+        AttributeError: When the reader class doesn't support file operations or _extension
+        IOError: When no files found locally or in object store
+    """
+    # Step 1: Check if files exist locally
+    local_files: List[str] = find_local_files_by_extension(
+        path, file_extension, file_names
+    )
+    if local_files:
+        logger.info(
+            f"Found {len(local_files)} {file_extension} files locally at: {path}"
+        )
+        return local_files
+
+    # Step 2: Try to download from object store
+    logger.info(
+        f"No local {file_extension} files found at {path}, checking object store..."
+    )
+
+    try:
+        # Determine what to download based on path type and filters
+        downloaded_paths: List[str] = []
+
+        if path.endswith(file_extension):
+            # Single file case (file_names validation already ensures this is valid)
+            source_path = get_object_store_prefix(path)
+            destination_path = os.path.join(TEMPORARY_PATH, source_path)
+            await ObjectStore.download_file(
+                source=source_path,
+                destination=destination_path,
+            )
+            downloaded_paths.append(destination_path)
+
+        elif file_names:
+            # Directory with specific files - download each file individually
+            for file_name in file_names:
+                file_path = os.path.join(path, file_name)
+                source_path = get_object_store_prefix(file_path)
+                destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                await ObjectStore.download_file(
+                    source=source_path,
+                    destination=destination_path,
+                )
+                downloaded_paths.append(destination_path)
+        else:
+            # Download entire directory
+            source_path = get_object_store_prefix(path)
+            destination_path = os.path.join(TEMPORARY_PATH, source_path)
+            await ObjectStore.download_prefix(
+                source=source_path,
+                destination=destination_path,
+            )
+            # Find the actual files in the downloaded directory
+            found_files = find_local_files_by_extension(
+                destination_path, file_extension, file_names
+            )
+            downloaded_paths.extend(found_files)
+
+        # Check results
+        if downloaded_paths:
+            logger.info(
+                f"Successfully downloaded {len(downloaded_paths)} {file_extension} files from object store"
+            )
+            return downloaded_paths
+        else:
+            raise IOError(
+                f"{IOError.OBJECT_STORE_READ_ERROR}: Downloaded from object store but no {file_extension} files found"
+            )
+
+    except Exception as e:
+        logger.error(f"Failed to download from object store: {str(e)}")
+        raise IOError(
+            f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: No {file_extension} files found locally at '{path}' and failed to download from object store. "
+            f"Error: {str(e)}"
+        )
+
+
+def estimate_dataframe_record_size(
+    dataframe: "pd.DataFrame", file_extension: str
+) -> int:
+    """Estimate File size of a DataFrame by sampling a few records.
+
+    Args:
+        dataframe (pd.DataFrame): The DataFrame to estimate the size of.
+        file_extension (str): The extension of the file to estimate the size of.
+
+    Returns:
+        int: The estimated size of the DataFrame in bytes.
+    """
+    if len(dataframe) == 0:
+        return 0
+
+    # Sample up to 10 records to estimate average size
+    sample_size = min(10, len(dataframe))
+    sample = dataframe.head(sample_size)
+    compression_factor = 1
+    if file_extension == JSON_FILE_EXTENSION:
+        sample_file = sample.to_json(orient="records", lines=True)
+    elif file_extension == PARQUET_FILE_EXTENSION:
+        sample_file = sample.to_parquet(index=False, compression="snappy")
+        compression_factor = 0.01
+    else:
+        raise ValueError(f"Unsupported file extension: {file_extension}")
+
+    if sample_file is not None:
+        avg_record_size = len(sample_file) / sample_size * compression_factor
+        return int(avg_record_size)
+
+    return 0
+
+
+def path_gen(
+    chunk_count: Optional[int] = None,
+    chunk_part: int = 0,
+    start_marker: Optional[str] = None,
+    end_marker: Optional[str] = None,
+    extension: str = ".json",
+) -> str:
+    """Generate a file path for a chunk.
+
+    Args:
+        chunk_start (Optional[int]): Starting index of the chunk, or None for single chunk.
+        chunk_count (int): Total number of chunks.
+        start_marker (Optional[str]): Start marker for query extraction.
+        end_marker (Optional[str]): End marker for query extraction.
+
+    Returns:
+        str: Generated file path for the chunk.
+    """
+    # For Query Extraction - use start and end markers without chunk count
+    if start_marker and end_marker:
+        return f"{start_marker}_{end_marker}{extension}"
+
+    # For regular chunking - include chunk count
+    if chunk_count is None:
+        return f"{str(chunk_part)}{extension}"
+    else:
+        return f"chunk-{str(chunk_count)}-part{str(chunk_part)}{extension}"
+
+
+def process_null_fields(
+    obj: Any,
+    preserve_fields: Optional[List[str]] = None,
+    null_to_empty_dict_fields: Optional[List[str]] = None,
+) -> Any:
+    """
+    By default the method removes null values from dictionaries and lists.
+    Except for the fields specified in preserve_fields.
+    And fields in null_to_empty_dict_fields are replaced with empty dict if null.
+
+    Args:
+        obj: The object to clean (dict, list, or other value)
+        preserve_fields: Optional list of field names that should be preserved even if they contain null values
+        null_to_empty_dict_fields: Optional list of field names that should be replaced with empty dict if null
+
+    Returns:
+        The cleaned object with null values removed
+    """
+    if isinstance(obj, dict):
+        result = {}
+        for k, v in obj.items():
+            # Handle null fields that should be converted to empty dicts
+            if k in (null_to_empty_dict_fields or []) and v is None:
+                result[k] = {}
+                continue
+
+            # Process the value recursively
+            processed_value = process_null_fields(
+                v, preserve_fields, null_to_empty_dict_fields
+            )
+
+            # Keep the field if it's in preserve_fields or has a non-None processed value
+            if k in (preserve_fields or []) or processed_value is not None:
+                result[k] = processed_value
+
+        return result
+    return obj
+
+
+def convert_datetime_to_epoch(data: Any) -> Any:
+    """Convert datetime objects to epoch timestamps in milliseconds.
+
+    Args:
+        data: The data to convert
+
+    Returns:
+        The converted data with datetime fields as epoch timestamps
+    """
+    if isinstance(data, datetime):
+        return int(data.timestamp() * 1000)
+    elif isinstance(data, dict):
+        return {k: convert_datetime_to_epoch(v) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [convert_datetime_to_epoch(item) for item in data]
+    return data
+
+
+def is_empty_dataframe(dataframe: Union["pd.DataFrame", "daft.DataFrame"]) -> bool:  # noqa: F821
+    """Check if a DataFrame is empty.
+
+    This function determines whether a DataFrame has any rows, supporting both
+    pandas and daft DataFrame types. For pandas DataFrames, it uses the `empty`
+    property, and for daft DataFrames, it checks if the row count is 0.
+
+    Args:
+        dataframe (Union[pd.DataFrame, daft.DataFrame]): The DataFrame to check,
+            can be either a pandas DataFrame or a daft DataFrame.
+
+    Returns:
+        bool: True if the DataFrame has no rows, False otherwise.
+
+    Note:
+        If daft is not available and a daft DataFrame is passed, the function
+        will log a warning and return True.
+    """
+    import pandas as pd
+
+    if isinstance(dataframe, pd.DataFrame):
+        return dataframe.empty
+
+    try:
+        import daft
+
+        if isinstance(dataframe, daft.DataFrame):
+            return dataframe.count_rows() == 0
+    except Exception:
+        logger.warning("Module daft not found")
+    return True

application_sdk/observability/observability.py

@@ -363,14 +363,14 @@ class AtlanObservability(Generic[T], ABC):
             logging.error(f"Error buffering log: {e}")
 
     async def _flush_records(self, records: List[Dict[str, Any]]):
-        """Flush records to parquet file and object store using ParquetOutput abstraction.
+        """Flush records to parquet file and object store using ParquetFileWriter.
 
         Args:
             records: List of records to flush
 
         This method:
         - Groups records by partition (year/month/day)
-        - Uses ParquetOutput abstraction for efficient writing
+        - Uses ParquetFileWriter for efficient writing
         - Automatically handles chunking, compression, and dual upload
         - Provides robust error handling per partition
         - Cleans up old records if enabled
@@ -395,7 +395,7 @@ class AtlanObservability(Generic[T], ABC):
                     partition_records[partition_path] = []
                 partition_records[partition_path].append(record)
 
-            # Write records to each partition using ParquetOutput abstraction
+                # Write records to each partition using ParquetFileWriter
             for partition_path, partition_data in partition_records.items():
                 # Create new dataframe from current records
                 new_df = pd.DataFrame(partition_data)
@@ -412,23 +412,27 @@ class AtlanObservability(Generic[T], ABC):
                     elif part.startswith("day="):
                         new_df["day"] = int(part.split("=")[1])
 
-                # Use new data directly - let ParquetOutput handle consolidation and merging
+                # Use new data directly - let ParquetFileWriter handle consolidation and merging
                 df = new_df
 
-                # Use ParquetOutput abstraction for efficient writing and uploading
+                # Use ParquetFileWriter for efficient writing and uploading
                 # Set the output path for this partition
                 try:
-                    # Lazy import and instantiation of ParquetOutput
-                    from application_sdk.outputs.parquet import ParquetOutput
+                    # Lazy import and instantiation of ParquetFileWriter
+                    from application_sdk.io.parquet import ParquetFileWriter
 
-                    parquet_output = ParquetOutput(
-                        output_path=partition_path,
+                    parquet_writer = ParquetFileWriter(
+                        path=partition_path,
                         chunk_start=0,
                         chunk_part=int(time()),
                     )
-                    await parquet_output.write_dataframe(dataframe=df)
-                except Exception as e:
-                    print(f"Error writing records to partition: {str(e)}")
+
+                    await parquet_writer._write_dataframe(dataframe=df)
+
+                except Exception as partition_error:
+                    logging.error(
+                        f"Error processing partition {partition_path}: {str(partition_error)}"
+                    )
 
             # Clean up old records if enabled
             if self._cleanup_enabled:

application_sdk/server/fastapi/middleware/logmiddleware.py

@@ -8,6 +8,7 @@ from starlette.types import ASGIApp
 
 from application_sdk.observability.context import request_context
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.server.fastapi.utils import EXCLUDED_LOG_PATHS
 
 logger = get_logger(__name__)
 
@@ -29,31 +30,36 @@ class LogMiddleware(BaseHTTPMiddleware):
         token = request_context.set({"request_id": request_id})
         start_time = time.time()
 
-        self.logger.info(
-            f"Request started for {request.method} {request.url.path}",
-            extra={
-                "method": request.method,
-                "path": request.url.path,
-                "request_id": request_id,
-                "url": str(request.url),
-                "client_host": request.client.host if request.client else None,
-            },
-        )
-
-        try:
-            response = await call_next(request)
-            duration = time.time() - start_time
+        # Skip logging for health check endpoints
+        should_log = request.url.path not in EXCLUDED_LOG_PATHS
 
+        if should_log:
             self.logger.info(
-                f"Request completed for {request.method} {request.url.path} {response.status_code}",
+                f"Request started for {request.method} {request.url.path}",
                 extra={
                     "method": request.method,
                     "path": request.url.path,
-                    "status_code": response.status_code,
-                    "duration_ms": round(duration * 1000, 2),
                     "request_id": request_id,
+                    "url": str(request.url),
+                    "client_host": request.client.host if request.client else None,
                 },
             )
+
+        try:
+            response = await call_next(request)
+            duration = time.time() - start_time
+
+            if should_log:
+                self.logger.info(
+                    f"Request completed for {request.method} {request.url.path} {response.status_code}",
+                    extra={
+                        "method": request.method,
+                        "path": request.url.path,
+                        "status_code": response.status_code,
+                        "duration_ms": round(duration * 1000, 2),
+                        "request_id": request_id,
+                    },
+                )
             return response
 
         except Exception as e:

application_sdk/server/fastapi/middleware/metrics.py

@@ -4,6 +4,7 @@ from fastapi import Request, Response
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
+from application_sdk.server.fastapi.utils import EXCLUDED_LOG_PATHS
 
 metrics = get_metrics()
 
@@ -24,29 +25,31 @@ class MetricsMiddleware(BaseHTTPMiddleware):
             method = request.method
             status_code = response.status_code
 
-            labels = {
-                "path": path,
-                "method": method,
-                "status": str(status_code),
-            }
-
-            # Record request count
-            metrics.record_metric(
-                name="http_requests_total",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels=labels,
-                description="Total number of HTTP requests",
-            )
-
-            # Record request latency
-            metrics.record_metric(
-                name="http_request_duration_seconds",
-                value=process_time,
-                metric_type=MetricType.HISTOGRAM,
-                labels=labels,
-                description="Duration of HTTP requests",
-                unit="seconds",
-            )
+            # Skip metrics for health check endpoints
+            if path not in EXCLUDED_LOG_PATHS:
+                labels = {
+                    "path": path,
+                    "method": method,
+                    "status": str(status_code),
+                }
+
+                # Record request count
+                metrics.record_metric(
+                    name="http_requests_total",
+                    value=1,
+                    metric_type=MetricType.COUNTER,
+                    labels=labels,
+                    description="Total number of HTTP requests",
+                )
+
+                # Record request latency
+                metrics.record_metric(
+                    name="http_request_duration_seconds",
+                    value=process_time,
+                    metric_type=MetricType.HISTOGRAM,
+                    labels=labels,
+                    description="Duration of HTTP requests",
+                    unit="seconds",
+                )
 
         return response

application_sdk/server/fastapi/models.py

@@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional, Type
 
 from pydantic import BaseModel, Field, RootModel
 
-from application_sdk.events.models import Event, EventFilter
+from application_sdk.interceptors.models import Event, EventFilter
 from application_sdk.workflows import WorkflowInterface
 
 

application_sdk/server/fastapi/routers/server.py

@@ -68,7 +68,7 @@ async def health():
         "processor": platform.processor(),
         "ram": str(round(psutil.virtual_memory().total / (1024.0**3))) + " GB",
     }
-    logger.info("Health check passed")
+    logger.debug("Health check passed")
     return info
 
 

application_sdk/server/fastapi/utils.py

@@ -7,6 +7,16 @@ error handlers and response formatters.
 from fastapi import status
 from fastapi.responses import JSONResponse
 
+# Paths to exclude from logging and metrics (health checks and event ingress)
+EXCLUDED_LOG_PATHS: frozenset[str] = frozenset(
+    {
+        "/server/health",
+        "/server/ready",
+        "/api/eventingress/",
+        "/api/eventingress",
+    }
+)
+
 
 def internal_server_error_handler(_, exc: Exception) -> JSONResponse:
     """Handle internal server errors in FastAPI applications.

application_sdk/services/eventstore.py

@@ -10,14 +10,14 @@ from datetime import datetime
 from dapr import clients
 from temporalio import activity, workflow
 
-from application_sdk.common.dapr_utils import is_component_registered
 from application_sdk.constants import (
     APPLICATION_NAME,
     DAPR_BINDING_OPERATION_CREATE,
     EVENT_STORE_NAME,
 )
-from application_sdk.events.models import Event, EventMetadata, WorkflowStates
+from application_sdk.interceptors.models import Event, EventMetadata, WorkflowStates
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services._utils import is_component_registered
 
 logger = get_logger(__name__)
 activity.logger = logger
@@ -47,7 +47,7 @@ class EventStore:
             a Temporal workflow or activity context.
 
         Examples:
-            >>> from application_sdk.events.models import Event
+            >>> from application_sdk.interceptors.models import Event
 
             >>> # Create basic event
             >>> event = Event(event_type="data.processed", data={"count": 100})
@@ -109,7 +109,7 @@ class EventStore:
             Exception: If there's an error during event publishing (logged but not re-raised).
 
         Examples:
-            >>> from application_sdk.events.models import Event
+            >>> from application_sdk.interceptors.models import Event
 
             >>> # Publish workflow status event
             >>> status_event = Event(

application_sdk/services/objectstore.py

@@ -459,9 +459,22 @@ class ObjectStore:
 
             logger.info(f"Found {len(file_list)} files to download from: {source}")
 
+            # Normalize source prefix to use forward slashes for comparison
+            normalized_source = cls._normalize_object_store_key(source)
+
             # Download each file
             for file_path in file_list:
-                local_file_path = os.path.join(destination, file_path)
+                normalized_file_path = cls._normalize_object_store_key(file_path)
+                if normalized_file_path.startswith(normalized_source):
+                    # Extract relative path after the prefix
+                    relative_path = normalized_file_path[
+                        len(normalized_source) :
+                    ].lstrip("/")
+                else:
+                    # Fallback to just the filename
+                    relative_path = os.path.basename(normalized_file_path)
+
+                local_file_path = os.path.join(destination, relative_path)
                 await cls.download_file(file_path, local_file_path, store_name)
 
             logger.info(f"Successfully downloaded all files from: {source}")

application_sdk/services/secretstore.py

@@ -20,7 +20,6 @@ from typing import Any, Dict
 
 from dapr.clients import DaprClient
 
-from application_sdk.common.dapr_utils import is_component_registered
 from application_sdk.common.error_codes import CommonError
 from application_sdk.constants import (
     DEPLOYMENT_NAME,
@@ -30,6 +29,7 @@ from application_sdk.constants import (
     SECRET_STORE_NAME,
 )
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services._utils import is_component_registered
 from application_sdk.services.statestore import StateStore, StateType
 
 logger = get_logger(__name__)

application_sdk/test_utils/hypothesis/strategies/outputs/json_output.py

@@ -62,7 +62,6 @@ def dataframe_strategy(draw) -> pd.DataFrame:
 json_output_config_strategy = st.fixed_dictionaries(
     {
         "output_path": safe_path_strategy,
-        "output_suffix": st.builds(lambda x: f"/{x}", safe_path_strategy),
         "output_prefix": output_prefix_strategy,
         "chunk_size": chunk_size_strategy,
     }

application_sdk/test_utils/hypothesis/strategies/server/fastapi/__init__.py

@@ -2,7 +2,7 @@ import json
 
 from hypothesis import strategies as st
 
-from application_sdk.events.models import Event
+from application_sdk.interceptors.models import Event
 
 # Strategy for generating auth credentials
 auth_credentials_strategy = st.fixed_dictionaries(

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "1.1.1"
+__version__ = "2.1.0"

application_sdk/worker.py

@@ -15,7 +15,7 @@ from temporalio.worker import Worker as TemporalWorker
 
 from application_sdk.clients.workflow import WorkflowClient
 from application_sdk.constants import DEPLOYMENT_NAME, MAX_CONCURRENT_ACTIVITIES
-from application_sdk.events.models import (
+from application_sdk.interceptors.models import (
     ApplicationEventNames,
     Event,
     EventTypes,