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

application_sdk/clients/temporal.py

@@ -151,7 +151,7 @@ class TemporalWorkflowClient(WorkflowClient):
                 await asyncio.sleep(refresh_interval)
 
                 # Get fresh token
-                token = await self.auth_manager.get_access_token()
+                token = await self.auth_manager.get_access_token(force_refresh=True)
                 if self.client:
                     self.client.api_key = token
                 logger.info("Updated client RPC metadata with fresh token")

application_sdk/outputs/parquet.py

@@ -1,5 +1,6 @@
 import os
-from typing import TYPE_CHECKING, List, Literal, Optional, Union
+from enum import Enum
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from temporalio import activity
 
@@ -18,6 +19,14 @@ if TYPE_CHECKING:
     import pandas as pd
 
 
+class WriteMode(Enum):
+    """Enumeration of write modes for Parquet output operations."""
+
+    APPEND = "append"
+    OVERWRITE = "overwrite"
+    OVERWRITE_PARTITIONS = "overwrite-partitions"
+
+
 class ParquetOutput(Output):
     """Output handler for writing data to Parquet files.
 
@@ -29,7 +38,6 @@ class ParquetOutput(Output):
         output_prefix (str): Prefix for files when uploading to object store.
         output_suffix (str): Suffix for output files.
         typename (Optional[str]): Type name of the entity e.g database, schema, table.
-        mode (str): Write mode for parquet files ("append" or "overwrite").
         chunk_size (int): Maximum number of records per chunk.
         total_record_count (int): Total number of records processed.
         chunk_count (int): Number of chunks created.
@@ -45,7 +53,6 @@ class ParquetOutput(Output):
         output_suffix: str = "",
         output_prefix: str = "",
         typename: Optional[str] = None,
-        write_mode: Literal["append", "overwrite", "overwrite-partitions"] = "append",
         chunk_size: Optional[int] = 100000,
         buffer_size: Optional[int] = 100000,
         total_record_count: int = 0,
@@ -61,7 +68,6 @@ class ParquetOutput(Output):
             output_suffix (str): Suffix for output files.
             output_prefix (str): Prefix for files when uploading to object store.
             typename (Optional[str], optional): Type name of the entity e.g database, schema, table.
-            mode (str, optional): Write mode for parquet files. Defaults to "append".
             chunk_size (int, optional): Maximum records per chunk. Defaults to 100000.
             total_record_count (int, optional): Initial total record count. Defaults to 0.
             chunk_count (int, optional): Initial chunk count. Defaults to 0.
@@ -78,7 +84,6 @@ class ParquetOutput(Output):
         self.output_suffix = output_suffix
         self.output_prefix = output_prefix
         self.typename = typename
-        self.write_mode = write_mode
         self.chunk_size = chunk_size
         self.buffer_size = buffer_size
         self.buffer: List[Union["pd.DataFrame", "daft.DataFrame"]] = []  # noqa: F821
@@ -103,7 +108,7 @@ class ParquetOutput(Output):
 
     def path_gen(
         self,
-        chunk_start: int | None = None,
+        chunk_start: Optional[int] = None,
         chunk_count: int = 0,
         start_marker: Optional[str] = None,
         end_marker: Optional[str] = None,
@@ -111,7 +116,7 @@ class ParquetOutput(Output):
         """Generate a file path for a chunk.
 
         Args:
-            chunk_start (int | None): Starting index of the chunk, or None for single chunk.
+            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.
@@ -182,7 +187,7 @@ class ParquetOutput(Output):
                 name="parquet_write_records",
                 value=len(dataframe),
                 metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": self.write_mode},
+                labels={"type": "pandas", "mode": WriteMode.APPEND.value},
                 description="Number of records written to Parquet files from pandas DataFrame",
             )
 
@@ -191,7 +196,7 @@ class ParquetOutput(Output):
                 name="parquet_chunks_written",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": self.write_mode},
+                labels={"type": "pandas", "mode": WriteMode.APPEND.value},
                 description="Number of chunks written to Parquet files",
             )
 
@@ -203,69 +208,109 @@ class ParquetOutput(Output):
                 name="parquet_write_errors",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": self.write_mode, "error": str(e)},
+                labels={
+                    "type": "pandas",
+                    "mode": WriteMode.APPEND.value,
+                    "error": str(e),
+                },
                 description="Number of errors while writing to Parquet files",
             )
             logger.error(f"Error writing pandas dataframe to parquet: {str(e)}")
             raise
 
-    async def write_daft_dataframe(self, dataframe: "daft.DataFrame"):  # noqa: F821
+    async def write_daft_dataframe(
+        self,
+        dataframe: "daft.DataFrame",  # noqa: F821
+        partition_cols: Optional[List] = None,
+        write_mode: Union[WriteMode, str] = WriteMode.APPEND,
+        morsel_size: int = 100_000,
+    ):
         """Write a daft DataFrame to Parquet files and upload to object store.
 
+        Uses Daft's native file size management to automatically split large DataFrames
+        into multiple parquet files based on the configured target file size. Supports
+        Hive partitioning for efficient data organization.
+
         Args:
             dataframe (daft.DataFrame): The DataFrame to write.
+            partition_cols (Optional[List]): Column names or expressions to use for Hive partitioning.
+                Can be strings (column names) or daft column expressions. If None (default), no partitioning is applied.
+            write_mode (Union[WriteMode, str]): Write mode for parquet files.
+                Use WriteMode.APPEND, WriteMode.OVERWRITE, WriteMode.OVERWRITE_PARTITIONS, or their string equivalents.
+            morsel_size (int): Default number of rows in a morsel used for the new local executor, when running locally on just a single machine,
+                Daft does not use partitions. Instead of using partitioning to control parallelism, the local execution engine performs a streaming-based
+                execution on small "morsels" of data, which provides much more stable memory utilization while improving the user experience with not having
+                to worry about partitioning.
+
+        Note:
+            - Daft automatically handles file chunking based on parquet_target_filesize
+            - Multiple files will be created if DataFrame exceeds DAPR limit
+            - If partition_cols is set, creates Hive-style directory structure
         """
         try:
+            import daft
+
+            # Convert string to enum if needed for backward compatibility
+            if isinstance(write_mode, str):
+                write_mode = WriteMode(write_mode)
+
             row_count = dataframe.count_rows()
             if row_count == 0:
                 return
 
+            # Use Daft's execution context for temporary configuration
+            with daft.execution_config_ctx(
+                parquet_target_filesize=self.max_file_size_bytes,
+                default_morsel_size=morsel_size,
+            ):
+                # Daft automatically handles file splitting and naming
+                dataframe.write_parquet(
+                    root_dir=self.output_path,
+                    write_mode=write_mode.value,
+                    partition_cols=partition_cols if partition_cols else [],
+                )
+
             # Update counters
             self.chunk_count += 1
             self.total_record_count += row_count
 
-            # Generate file path using path_gen function
-            if self.start_marker and self.end_marker:
-                file_path = self.output_path
-            else:
-                file_path = f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count, self.start_marker, self.end_marker)}"
-
-            # Write the dataframe to parquet using daft
-            dataframe.write_parquet(
-                file_path,
-                write_mode=self.write_mode,
-            )
-
             # Record metrics for successful write
             self.metrics.record_metric(
                 name="parquet_write_records",
                 value=row_count,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "daft", "mode": self.write_mode},
+                labels={"type": "daft", "mode": write_mode.value},
                 description="Number of records written to Parquet files from daft DataFrame",
             )
 
-            # Record chunk metrics
+            # Record operation metrics (note: actual file count may be higher due to Daft's splitting)
             self.metrics.record_metric(
-                name="parquet_chunks_written",
+                name="parquet_write_operations",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "daft", "mode": self.write_mode},
-                description="Number of chunks written to Parquet files",
+                labels={"type": "daft", "mode": write_mode.value},
+                description="Number of write operations to Parquet files",
             )
 
-            # Upload the file to object store
-            await ObjectStore.upload_file(
-                source=file_path,
-                destination=get_object_store_prefix(file_path),
+            #  Upload the entire directory (contains multiple parquet files created by Daft)
+            if write_mode == WriteMode.OVERWRITE:
+                # Delete the directory from object store
+                await ObjectStore.delete_prefix(
+                    prefix=get_object_store_prefix(self.output_path)
+                )
+
+            await ObjectStore.upload_prefix(
+                source=self.output_path,
+                destination=get_object_store_prefix(self.output_path),
             )
+
         except Exception as e:
             # Record metrics for failed write
             self.metrics.record_metric(
                 name="parquet_write_errors",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "daft", "mode": self.write_mode, "error": str(e)},
+                labels={"type": "daft", "mode": write_mode, "error": str(e)},
                 description="Number of errors while writing to Parquet files",
             )
             logger.error(f"Error writing daft dataframe to parquet: {str(e)}")
@@ -279,7 +324,7 @@ class ParquetOutput(Output):
         """
         return self.output_path
 
-    async def _flush_buffer(self, chunk_part):
+    async def _flush_buffer(self, chunk_part: int):
         """Flush the current buffer to a Parquet file.
 
         This method combines all DataFrames in the buffer, writes them to a Parquet file,

application_sdk/services/objectstore.py

@@ -26,6 +26,31 @@ class ObjectStore:
     OBJECT_CREATE_OPERATION = "create"
     OBJECT_GET_OPERATION = "get"
     OBJECT_LIST_OPERATION = "list"
+    OBJECT_DELETE_OPERATION = "delete"
+
+    @classmethod
+    def _create_file_metadata(cls, key: str) -> dict[str, str]:
+        """Create metadata for file operations (get, delete, create).
+
+        Args:
+            key: The file key/path.
+
+        Returns:
+            Metadata dictionary with key, fileName, and blobName fields.
+        """
+        return {"key": key, "fileName": key, "blobName": key}
+
+    @classmethod
+    def _create_list_metadata(cls, prefix: str) -> dict[str, str]:
+        """Create metadata for list operations.
+
+        Args:
+            prefix: The prefix to list files under.
+
+        Returns:
+            Metadata dictionary with prefix and fileName fields, or empty dict if no prefix.
+        """
+        return {"prefix": prefix, "fileName": prefix} if prefix else {}
 
     @classmethod
     async def list_files(
@@ -44,12 +69,11 @@ class ObjectStore:
             Exception: If there's an error listing files from the object store.
         """
         try:
-            metadata = {"prefix": prefix, "fileName": prefix} if prefix else {}
             data = json.dumps({"prefix": prefix}).encode("utf-8") if prefix else ""
 
             response_data = await cls._invoke_dapr_binding(
                 operation=cls.OBJECT_LIST_OPERATION,
-                metadata=metadata,
+                metadata=cls._create_list_metadata(prefix),
                 data=data,
                 store_name=store_name,
             )
@@ -105,12 +129,11 @@ class ObjectStore:
             Exception: If there's an error getting the file from the object store.
         """
         try:
-            metadata = {"key": key, "fileName": key, "blobName": key}
             data = json.dumps({"key": key}).encode("utf-8") if key else ""
 
             response_data = await cls._invoke_dapr_binding(
                 operation=cls.OBJECT_GET_OPERATION,
-                metadata=metadata,
+                metadata=cls._create_file_metadata(key),
                 data=data,
                 store_name=store_name,
             )
@@ -144,20 +167,68 @@ class ObjectStore:
             return False
 
     @classmethod
-    async def delete(
+    async def delete_file(
         cls, key: str, store_name: str = DEPLOYMENT_OBJECT_STORE_NAME
     ) -> None:
-        """Delete a file or all files under a prefix from the object store.
+        """Delete a single file from the object store.
 
         Args:
-            key: The file path or prefix to delete.
+            key: The file path to delete.
             store_name: Name of the Dapr object store binding to use.
 
-        Note:
-            This method is not implemented as it's not commonly used in the current codebase.
-            Can be implemented when needed based on the underlying object store capabilities.
+        Raises:
+            Exception: If there's an error deleting the file from the object store.
+        """
+        try:
+            data = json.dumps({"key": key}).encode("utf-8")
+
+            await cls._invoke_dapr_binding(
+                operation=cls.OBJECT_DELETE_OPERATION,
+                metadata=cls._create_file_metadata(key),
+                data=data,
+                store_name=store_name,
+            )
+            logger.debug(f"Successfully deleted file: {key}")
+        except Exception as e:
+            logger.error(f"Error deleting file {key}: {str(e)}")
+            raise
+
+    @classmethod
+    async def delete_prefix(
+        cls, prefix: str, store_name: str = DEPLOYMENT_OBJECT_STORE_NAME
+    ) -> None:
+        """Delete all files under a prefix from the object store.
+
+        Args:
+            prefix: The prefix path to delete all files under.
+            store_name: Name of the Dapr object store binding to use.
+
+        Raises:
+            Exception: If there's an error deleting files from the object store.
         """
-        raise NotImplementedError("Delete operation not yet implemented")
+        try:
+            # First, list all files under the prefix
+            files_to_delete = await cls.list_files(prefix=prefix, store_name=store_name)
+
+            if not files_to_delete:
+                logger.info(f"No files found under prefix: {prefix}")
+                return
+
+            logger.info(f"Deleting {len(files_to_delete)} files under prefix: {prefix}")
+
+            # Delete each file individually
+            for file_path in files_to_delete:
+                try:
+                    await cls.delete_file(key=file_path, store_name=store_name)
+                except Exception as e:
+                    logger.warning(f"Failed to delete file {file_path}: {str(e)}")
+                    # Continue with other files even if one fails
+
+            logger.info(f"Successfully deleted all files under prefix: {prefix}")
+
+        except Exception as e:
+            logger.error(f"Error deleting files under prefix {prefix}: {str(e)}")
+            raise
 
     @classmethod
     async def upload_file(
@@ -191,17 +262,11 @@ class ObjectStore:
             logger.error(f"Error reading file {source}: {str(e)}")
             raise e
 
-        metadata = {
-            "key": destination,
-            "blobName": destination,
-            "fileName": destination,
-        }
-
         try:
             await cls._invoke_dapr_binding(
                 operation=cls.OBJECT_CREATE_OPERATION,
                 data=file_content,
-                metadata=metadata,
+                metadata=cls._create_file_metadata(destination),
                 store_name=store_name,
             )
             logger.debug(f"Successfully uploaded file: {destination}")

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "0.1.1rc38"
+__version__ = "0.1.1rc39"

{atlan_application_sdk-0.1.1rc38.dist-info → atlan_application_sdk-0.1.1rc39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc38
+Version: 0.1.1rc39
 Summary: Atlan Application SDK is a Python library for developing applications on the Atlan Platform
 Project-URL: Repository, https://github.com/atlanhq/application-sdk
 Project-URL: Documentation, https://github.com/atlanhq/application-sdk/README.md

{atlan_application_sdk-0.1.1rc38.dist-info → atlan_application_sdk-0.1.1rc39.dist-info}/RECORD

@@ -1,6 +1,6 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
 application_sdk/constants.py,sha256=GzwZO0pa9M-FgibmfIs1lh-Fwo06K9Tk6WzGqMyJgpI,10362
-application_sdk/version.py,sha256=4PQKkm_QOy4mEJgwY5eIcUeeEIpIWPjzc_pKiJQmpXw,88
+application_sdk/version.py,sha256=c9l4xzOeaiyxuYpkhXgaiBr52RJe_2lv_aUnX-NAeR4,88
 application_sdk/worker.py,sha256=i5f0AeKI39IfsLO05QkwC6uMz0zDPSJqP7B2byri1VI,7489
 application_sdk/activities/__init__.py,sha256=QaXLOBYbb0zPOY5kfDQh56qbXQFaYNXOjJ5PCvatiZ4,9530
 application_sdk/activities/lock_management.py,sha256=L__GZ9BsArwU1ntYwAgCKsSjCqN6QBeOfT-OT4WyD4Y,3983
@@ -21,7 +21,7 @@ application_sdk/clients/atlan_auth.py,sha256=D7FuNqv81ohNXLJtdx1AFw_jU6a3g0Pw614
 application_sdk/clients/base.py,sha256=TIn3pG89eXUc1XSYf4jk66m1vajWp0WxcCQOOltdazA,14021
 application_sdk/clients/redis.py,sha256=IfAD32vLp88BCvsDTaQtxFHxzHlEx4V7TK7h1HwDDBg,15917
 application_sdk/clients/sql.py,sha256=tW89SHuuWdU5jv8lDUP5AUCEpR2CF_5TyUvYDCBHses,17880
-application_sdk/clients/temporal.py,sha256=D5tr6i6DfkTQ--kZYU-sz7yOEsxigf-g3NDnwpOhR2I,18217
+application_sdk/clients/temporal.py,sha256=MEGNT1_crWAn-vdfcBUH0A7IKvKDDSAaiEpGCS7gas4,18235
 application_sdk/clients/utils.py,sha256=zLFOJbTr_6TOqnjfVFGY85OtIXZ4FQy_rquzjaydkbY,779
 application_sdk/clients/workflow.py,sha256=6bSqmA3sNCk9oY68dOjBUDZ9DhNKQxPD75qqE0cfldc,6104
 application_sdk/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -70,7 +70,7 @@ application_sdk/observability/decorators/observability_decorator.py,sha256=JNrWN
 application_sdk/outputs/__init__.py,sha256=HIENr2w9gu6u3sF_nvraj45yk53NDAddtaXSUHIVBjs,9469
 application_sdk/outputs/iceberg.py,sha256=IGtj5WDgqLu6vzDEvw5DLsKsjm29Krto3AHvWpemr0A,5311
 application_sdk/outputs/json.py,sha256=zyYQjGj5tb7bJhNt3ObwsuHT6Gakj8qNey-siUlWdP4,15065
-application_sdk/outputs/parquet.py,sha256=pJkOw-CV-JXr0Q4nJXCu0SScbv3I0usONZ4CT1KYxYI,13973
+application_sdk/outputs/parquet.py,sha256=AHcsuIVU0C-yNmqk4oUuDwdxCBSjaAxJwe2v7EklVsg,16220
 application_sdk/server/__init__.py,sha256=KTqE1YPw_3WDVMWatJUuf9OOiobLM2K5SMaBrI62sCo,1568
 application_sdk/server/fastapi/__init__.py,sha256=YOdWNE-qqiXfo-exvxPg8T0PSuOxTdeSetUn6-BXxZg,27704
 application_sdk/server/fastapi/models.py,sha256=K6eNl3XXiTXKUvRTpq3oqdGH3jY1-ApobXma04J86fE,6665
@@ -82,7 +82,7 @@ application_sdk/server/fastapi/routers/server.py,sha256=vfHQwZCysThzfeVFNVW1IjuA
 application_sdk/services/__init__.py,sha256=H-5HZEPdr53MUfAggyHqHhRXDRLZFZsxvJgWbr257Ds,465
 application_sdk/services/atlan_storage.py,sha256=TKzXxu0yXeUcmZehwp8PcnQTC4A9w9RlZ0Fl-Xp1bLE,8509
 application_sdk/services/eventstore.py,sha256=X03JzodKByXh8w8nOl658rnnZfMFTj0IkmiLVbd6IN8,6729
-application_sdk/services/objectstore.py,sha256=0TO7TYRvZlQHU3OH6PCA6tGO49E2lM7eCvV-fRWLiqU,14098
+application_sdk/services/objectstore.py,sha256=85e_68ubUsoj3riPSgPorFvJR0vnAlkSq3uglRjUtIA,16402
 application_sdk/services/secretstore.py,sha256=UpyLLcdMia1tqFCpRrn-lE9AnERAt2iGVzET6QqkmqI,13976
 application_sdk/services/statestore.py,sha256=CQuKq4FXPS0ebDH0e0cfTTAjvsIlrA1zz1MpsWCiWnM,9562
 application_sdk/test_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -141,8 +141,8 @@ application_sdk/workflows/metadata_extraction/__init__.py,sha256=jHUe_ZBQ66jx8bg
 application_sdk/workflows/metadata_extraction/sql.py,sha256=BhaZavEL8H3Jvf28FGcHtZwqdsUT_EHZ4VTqiaieWek,12278
 application_sdk/workflows/query_extraction/__init__.py,sha256=n066_CX5RpJz6DIxGMkKS3eGSRg03ilaCtsqfJWQb7Q,117
 application_sdk/workflows/query_extraction/sql.py,sha256=kT_JQkLCRZ44ZpaC4QvPL6DxnRIIVh8gYHLqRbMI-hA,4826
-atlan_application_sdk-0.1.1rc38.dist-info/METADATA,sha256=dQKtTnaVMPF0DU96vA850Gl6gzpRSmuL48ovbq1wD5o,5567
-atlan_application_sdk-0.1.1rc38.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-atlan_application_sdk-0.1.1rc38.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-atlan_application_sdk-0.1.1rc38.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
-atlan_application_sdk-0.1.1rc38.dist-info/RECORD,,
+atlan_application_sdk-0.1.1rc39.dist-info/METADATA,sha256=D1HHgG-YRM3VS43mQGsST_RJTkkDacnCSnFszqPfSL4,5567
+atlan_application_sdk-0.1.1rc39.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+atlan_application_sdk-0.1.1rc39.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-0.1.1rc39.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-0.1.1rc39.dist-info/RECORD,,