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

application_sdk/outputs/__init__.py

@@ -1,453 +0,0 @@
-"""Output module for handling data output operations.
-
-This module provides base classes and utilities for handling various types of data outputs
-in the application, including file outputs and object store interactions.
-"""
-
-import gc
-import inspect
-import os
-from abc import ABC, abstractmethod
-from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    AsyncGenerator,
-    Dict,
-    Generator,
-    List,
-    Optional,
-    Union,
-    cast,
-)
-
-import orjson
-from temporalio import activity
-
-from application_sdk.activities.common.models import ActivityStatistics
-from application_sdk.activities.common.utils import get_object_store_prefix
-from application_sdk.common.dataframe_utils import is_empty_dataframe
-from application_sdk.constants import ENABLE_ATLAN_UPLOAD, UPSTREAM_OBJECT_STORE_NAME
-from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.observability.metrics_adaptor import MetricType
-from application_sdk.services.objectstore import ObjectStore
-
-logger = get_logger(__name__)
-activity.logger = logger
-
-
-if TYPE_CHECKING:
-    import daft  # type: ignore
-    import pandas as pd
-
-
-class WriteMode(Enum):
-    """Enumeration of write modes for output operations."""
-
-    APPEND = "append"
-    OVERWRITE = "overwrite"
-    OVERWRITE_PARTITIONS = "overwrite-partitions"
-
-
-class Output(ABC):
-    """Abstract base class for output handlers.
-
-    This class defines the interface for output handlers that can write data
-    to various destinations in different formats.
-
-    Attributes:
-        output_path (str): Path where the output will be written.
-        upload_file_prefix (str): Prefix for files when uploading to object store.
-        total_record_count (int): Total number of records processed.
-        chunk_count (int): Number of chunks the output was split into.
-    """
-
-    output_path: str
-    output_prefix: str
-    total_record_count: int
-    chunk_count: int
-    chunk_part: int
-    buffer_size: int
-    max_file_size_bytes: int
-    current_buffer_size: int
-    current_buffer_size_bytes: int
-    partitions: List[int]
-
-    def estimate_dataframe_record_size(self, dataframe: "pd.DataFrame") -> int:
-        """Estimate File size of a DataFrame by sampling a few records."""
-        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)
-        file_type = type(self).__name__.lower().replace("output", "")
-        compression_factor = 1
-        if file_type == "json":
-            sample_file = sample.to_json(orient="records", lines=True)
-        else:
-            sample_file = sample.to_parquet(index=False, compression="snappy")
-            compression_factor = 0.01
-        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(
-        self,
-        chunk_count: Optional[int] = None,
-        chunk_part: int = 0,
-        start_marker: Optional[str] = None,
-        end_marker: Optional[str] = None,
-    ) -> 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}{self._EXTENSION}"
-
-        # For regular chunking - include chunk count
-        if chunk_count is None:
-            return f"{str(chunk_part)}{self._EXTENSION}"
-        else:
-            return f"chunk-{str(chunk_count)}-part{str(chunk_part)}{self._EXTENSION}"
-
-    def process_null_fields(
-        self,
-        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 = self.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
-
-    async def write_batched_dataframe(
-        self,
-        batched_dataframe: Union[
-            AsyncGenerator["pd.DataFrame", None], Generator["pd.DataFrame", None, None]
-        ],
-    ):
-        """Write a batched pandas DataFrame to Output.
-
-        This method writes the DataFrame to Output provided, potentially splitting it
-        into chunks based on chunk_size and buffer_size settings.
-
-        Args:
-            dataframe (pd.DataFrame): The DataFrame to write.
-
-        Note:
-            If the DataFrame is empty, the method returns without writing.
-        """
-        try:
-            if inspect.isasyncgen(batched_dataframe):
-                async for dataframe in batched_dataframe:
-                    if not is_empty_dataframe(dataframe):
-                        await self.write_dataframe(dataframe)
-            else:
-                # Cast to Generator since we've confirmed it's not an AsyncGenerator
-                sync_generator = cast(
-                    Generator["pd.DataFrame", None, None], batched_dataframe
-                )
-                for dataframe in sync_generator:
-                    if not is_empty_dataframe(dataframe):
-                        await self.write_dataframe(dataframe)
-        except Exception as e:
-            logger.error(f"Error writing batched dataframe: {str(e)}")
-            raise
-
-    async def write_dataframe(self, dataframe: "pd.DataFrame"):
-        """Write a pandas DataFrame to Parquet files and upload to object store.
-
-        Args:
-            dataframe (pd.DataFrame): The DataFrame to write.
-        """
-        try:
-            if self.chunk_start is None:
-                self.chunk_part = 0
-            if len(dataframe) == 0:
-                return
-
-            chunk_size_bytes = self.estimate_dataframe_record_size(dataframe)
-
-            for i in range(0, len(dataframe), self.buffer_size):
-                chunk = dataframe[i : i + self.buffer_size]
-
-                if (
-                    self.current_buffer_size_bytes + chunk_size_bytes
-                    > self.max_file_size_bytes
-                ):
-                    output_file_name = f"{self.output_path}/{self.path_gen(self.chunk_count, self.chunk_part)}"
-                    if os.path.exists(output_file_name):
-                        await self._upload_file(output_file_name)
-                        self.chunk_part += 1
-
-                self.current_buffer_size += len(chunk)
-                self.current_buffer_size_bytes += chunk_size_bytes * len(chunk)
-                await self._flush_buffer(chunk, self.chunk_part)
-
-                del chunk
-                gc.collect()
-
-            if self.current_buffer_size_bytes > 0:
-                # Finally upload the final file to the object store
-                output_file_name = f"{self.output_path}/{self.path_gen(self.chunk_count, self.chunk_part)}"
-                if os.path.exists(output_file_name):
-                    await self._upload_file(output_file_name)
-                    self.chunk_part += 1
-
-            # Record metrics for successful write
-            self.metrics.record_metric(
-                name="write_records",
-                value=len(dataframe),
-                metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": WriteMode.APPEND.value},
-                description="Number of records written to files from pandas DataFrame",
-            )
-
-            # Record chunk metrics
-            self.metrics.record_metric(
-                name="chunks_written",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": WriteMode.APPEND.value},
-                description="Number of chunks written to files",
-            )
-
-            # If chunk_start is set we don't want to increment the chunk_count
-            # Since it should only increment the chunk_part in this case
-            if self.chunk_start is None:
-                self.chunk_count += 1
-            self.partitions.append(self.chunk_part)
-        except Exception as e:
-            # Record metrics for failed write
-            self.metrics.record_metric(
-                name="write_errors",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={
-                    "type": "pandas",
-                    "mode": WriteMode.APPEND.value,
-                    "error": str(e),
-                },
-                description="Number of errors while writing to files",
-            )
-            logger.error(f"Error writing pandas dataframe to files: {str(e)}")
-            raise
-
-    async def write_batched_daft_dataframe(
-        self,
-        batched_dataframe: Union[
-            AsyncGenerator["daft.DataFrame", None],  # noqa: F821
-            Generator["daft.DataFrame", None, None],  # noqa: F821
-        ],
-    ):
-        """Write a batched daft DataFrame to JSON files.
-
-        This method writes the DataFrame to JSON files, potentially splitting it
-        into chunks based on chunk_size and buffer_size settings.
-
-        Args:
-            dataframe (daft.DataFrame): The DataFrame to write.
-
-        Note:
-            If the DataFrame is empty, the method returns without writing.
-        """
-        try:
-            if inspect.isasyncgen(batched_dataframe):
-                async for dataframe in batched_dataframe:
-                    if not is_empty_dataframe(dataframe):
-                        await self.write_daft_dataframe(dataframe)
-            else:
-                # Cast to Generator since we've confirmed it's not an AsyncGenerator
-                sync_generator = cast(
-                    Generator["daft.DataFrame", None, None], batched_dataframe
-                )  # noqa: F821
-                for dataframe in sync_generator:
-                    if not is_empty_dataframe(dataframe):
-                        await self.write_daft_dataframe(dataframe)
-        except Exception as e:
-            logger.error(f"Error writing batched daft dataframe: {str(e)}")
-
-    @abstractmethod
-    async def write_daft_dataframe(self, dataframe: "daft.DataFrame"):  # noqa: F821
-        """Write a daft DataFrame to the output destination.
-
-        Args:
-            dataframe (daft.DataFrame): The DataFrame to write.
-        """
-        pass
-
-    async def get_statistics(
-        self, typename: Optional[str] = None
-    ) -> ActivityStatistics:
-        """Returns statistics about the output.
-
-        This method returns a ActivityStatistics object with total record count and chunk count.
-
-        Args:
-            typename (str): Type name of the entity e.g database, schema, table.
-
-        Raises:
-            ValidationError: If the statistics data is invalid
-            Exception: If there's an error writing the statistics
-        """
-        try:
-            statistics = await self.write_statistics(typename)
-            if not statistics:
-                raise ValueError("No statistics data available")
-            statistics = ActivityStatistics.model_validate(statistics)
-            if typename:
-                statistics.typename = typename
-            return statistics
-        except Exception as e:
-            logger.error(f"Error getting statistics: {str(e)}")
-            raise
-
-    async def _upload_file(self, file_name: str):
-        """Upload a file to the object store."""
-        if ENABLE_ATLAN_UPLOAD:
-            await ObjectStore.upload_file(
-                source=file_name,
-                store_name=UPSTREAM_OBJECT_STORE_NAME,
-                retain_local_copy=True,
-                destination=get_object_store_prefix(file_name),
-            )
-        await ObjectStore.upload_file(
-            source=file_name,
-            destination=get_object_store_prefix(file_name),
-        )
-
-        self.current_buffer_size_bytes = 0
-
-    async def _flush_buffer(self, chunk: "pd.DataFrame", chunk_part: int):
-        """Flush the current buffer to a JSON file.
-
-        This method combines all DataFrames in the buffer, writes them to a JSON file,
-        and uploads the file to the object store.
-
-        Note:
-            If the buffer is empty or has no records, the method returns without writing.
-        """
-        try:
-            if not is_empty_dataframe(chunk):
-                self.total_record_count += len(chunk)
-                output_file_name = (
-                    f"{self.output_path}/{self.path_gen(self.chunk_count, chunk_part)}"
-                )
-                await self.write_chunk(chunk, output_file_name)
-
-                self.current_buffer_size = 0
-
-                # Record chunk metrics
-                self.metrics.record_metric(
-                    name="chunks_written",
-                    value=1,
-                    metric_type=MetricType.COUNTER,
-                    labels={"type": "output"},
-                    description="Number of chunks written to files",
-                )
-
-        except Exception as e:
-            # Record metrics for failed write
-            self.metrics.record_metric(
-                name="write_errors",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={"type": "output", "error": str(e)},
-                description="Number of errors while writing to files",
-            )
-            logger.error(f"Error flushing buffer to files: {str(e)}")
-            raise e
-
-    async def write_statistics(
-        self, typename: Optional[str] = None
-    ) -> Optional[Dict[str, Any]]:
-        """Write statistics about the output to a JSON file.
-
-        This method writes statistics including total record count and chunk count
-        to a JSON file and uploads it to the object store.
-
-        Raises:
-            Exception: If there's an error writing or uploading the statistics.
-        """
-        try:
-            # prepare the statistics
-            statistics = {
-                "total_record_count": self.total_record_count,
-                "chunk_count": len(self.partitions),
-                "partitions": self.partitions,
-            }
-
-            # Ensure typename is included in the statistics payload (if provided)
-            if typename:
-                statistics["typename"] = typename
-
-            # Write the statistics to a json file inside a dedicated statistics/ folder
-            statistics_dir = os.path.join(self.output_path, "statistics")
-            os.makedirs(statistics_dir, exist_ok=True)
-            output_file_name = os.path.join(statistics_dir, "statistics.json.ignore")
-            # If chunk_start is provided, include it in the statistics filename
-            try:
-                cs = getattr(self, "chunk_start", None)
-                if cs is not None:
-                    output_file_name = os.path.join(
-                        statistics_dir, f"statistics-chunk-{cs}.json.ignore"
-                    )
-            except Exception:
-                # If accessing chunk_start fails, fallback to default filename
-                pass
-
-            # Write the statistics dictionary to the JSON file
-            with open(output_file_name, "wb") as f:
-                f.write(orjson.dumps(statistics))
-
-            destination_file_path = get_object_store_prefix(output_file_name)
-            # Push the file to the object store
-            await ObjectStore.upload_file(
-                source=output_file_name,
-                destination=destination_file_path,
-            )
-
-            return statistics
-        except Exception as e:
-            logger.error(f"Error writing statistics: {str(e)}")

application_sdk/outputs/iceberg.py

@@ -1,139 +0,0 @@
-from typing import TYPE_CHECKING, Union
-
-from pyiceberg.catalog import Catalog
-from pyiceberg.table import Table
-from temporalio import activity
-
-from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
-from application_sdk.outputs import Output
-
-logger = get_logger(__name__)
-activity.logger = logger
-
-if TYPE_CHECKING:
-    import daft
-    import pandas as pd
-
-
-class IcebergOutput(Output):
-    """
-    Iceberg Output class to write data to Iceberg tables using daft and pandas
-    """
-
-    def __init__(
-        self,
-        iceberg_catalog: Catalog,
-        iceberg_namespace: str,
-        iceberg_table: Union[str, Table],
-        mode: str = "append",
-        total_record_count: int = 0,
-        chunk_count: int = 0,
-        retain_local_copy: bool = False,
-    ):
-        """Initialize the Iceberg output class.
-
-        Args:
-            iceberg_catalog (Catalog): Iceberg catalog object.
-            iceberg_namespace (str): Iceberg namespace.
-            iceberg_table (Union[str, Table]): Iceberg table object or table name.
-            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
-        self.iceberg_catalog = iceberg_catalog
-        self.iceberg_namespace = iceberg_namespace
-        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"):
-        """
-        Method to write the pandas dataframe to an iceberg table
-        """
-        try:
-            import daft
-
-            if len(dataframe) == 0:
-                return
-            # convert the pandas dataframe to a daft dataframe
-            daft_dataframe = daft.from_pandas(dataframe)
-            await self.write_daft_dataframe(daft_dataframe)
-
-            # Record metrics for successful write
-            self.metrics.record_metric(
-                name="iceberg_write_records",
-                value=len(dataframe),
-                metric_type=MetricType.COUNTER,
-                labels={"mode": self.mode, "type": "pandas"},
-                description="Number of records written to Iceberg table from pandas DataFrame",
-            )
-        except Exception as e:
-            # Record metrics for failed write
-            self.metrics.record_metric(
-                name="iceberg_write_errors",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={"mode": self.mode, "type": "pandas", "error": str(e)},
-                description="Number of errors while writing to Iceberg table",
-            )
-            logger.error(f"Error writing pandas dataframe to iceberg table: {str(e)}")
-            raise e
-
-    async def write_daft_dataframe(self, dataframe: "daft.DataFrame"):  # noqa: F821
-        """
-        Method to write the daft dataframe to an iceberg table
-        """
-        try:
-            if dataframe.count_rows() == 0:
-                return
-            # Create a new table in the iceberg catalog
-            self.chunk_count += 1
-            self.total_record_count += dataframe.count_rows()
-
-            # check if iceberg table is already created
-            if isinstance(self.iceberg_table, Table):
-                # if yes, use the existing iceberg table
-                table = self.iceberg_table
-            else:
-                # if not, create a new table in the iceberg catalog
-                table = self.iceberg_catalog.create_table_if_not_exists(
-                    f"{self.iceberg_namespace}.{self.iceberg_table}",
-                    schema=dataframe.to_arrow().schema,
-                )
-            # write the dataframe to the iceberg table
-            dataframe.write_iceberg(table, mode=self.mode)
-
-            # Record metrics for successful write
-            self.metrics.record_metric(
-                name="iceberg_write_records",
-                value=dataframe.count_rows(),
-                metric_type=MetricType.COUNTER,
-                labels={"mode": self.mode, "type": "daft"},
-                description="Number of records written to Iceberg table from daft DataFrame",
-            )
-
-            # Record chunk metrics
-            self.metrics.record_metric(
-                name="iceberg_chunks_written",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={"mode": self.mode},
-                description="Number of chunks written to Iceberg table",
-            )
-        except Exception as e:
-            # Record metrics for failed write
-            self.metrics.record_metric(
-                name="iceberg_write_errors",
-                value=1,
-                metric_type=MetricType.COUNTER,
-                labels={"mode": self.mode, "type": "daft", "error": str(e)},
-                description="Number of errors while writing to Iceberg table",
-            )
-            logger.error(f"Error writing daft dataframe to iceberg table: {str(e)}")
-            raise e