atlan-application-sdk 0.1.1rc35__py3-none-any.whl → 0.1.1rc37__py3-none-any.whl

application_sdk/interceptors/lock.py

@@ -0,0 +1,139 @@
+"""Redis lock interceptor for Temporal workflows.
+
+Manages distributed locks for activities decorated with @needs_lock using
+separate lock acquisition and release activities to avoid workflow deadlocks.
+"""
+
+from datetime import timedelta
+from typing import Any, Dict, Optional, Type
+
+from temporalio import workflow
+from temporalio.common import RetryPolicy
+from temporalio.worker import (
+    Interceptor,
+    StartActivityInput,
+    WorkflowInboundInterceptor,
+    WorkflowInterceptorClassInput,
+    WorkflowOutboundInterceptor,
+)
+
+from application_sdk.common.error_codes import WorkflowError
+from application_sdk.constants import (
+    APPLICATION_NAME,
+    IS_LOCKING_DISABLED,
+    LOCK_METADATA_KEY,
+)
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+class RedisLockInterceptor(Interceptor):
+    """Main interceptor class for Redis distributed locking."""
+
+    def __init__(self, activities: Dict[str, Any]):
+        """Initialize Redis lock interceptor.
+
+        Args:
+            activities: Dictionary mapping activity names to activity functions
+        """
+        self.activities = activities
+
+    def workflow_interceptor_class(
+        self, input: WorkflowInterceptorClassInput
+    ) -> Optional[Type[WorkflowInboundInterceptor]]:
+        activities = self.activities
+
+        class RedisLockWorkflowInboundInterceptor(WorkflowInboundInterceptor):
+            """Inbound interceptor that manages Redis locks for activities."""
+
+            def init(self, outbound: WorkflowOutboundInterceptor) -> None:
+                """Initialize with Redis lock outbound interceptor."""
+                lock_outbound = RedisLockOutboundInterceptor(outbound, activities)
+                super().init(lock_outbound)
+
+        return RedisLockWorkflowInboundInterceptor
+
+
+class RedisLockOutboundInterceptor(WorkflowOutboundInterceptor):
+    """Outbound interceptor that acquires Redis locks before activity execution."""
+
+    def __init__(self, next: WorkflowOutboundInterceptor, activities: Dict[str, Any]):
+        super().__init__(next)
+        self.activities = activities
+
+    async def start_activity(  # type: ignore[override]
+        self, input: StartActivityInput
+    ) -> workflow.ActivityHandle[Any]:
+        """Start activity with distributed lock if required."""
+
+        # Check if activity needs locking
+        activity_fn = self.activities.get(input.activity)
+        if (
+            not activity_fn
+            or not hasattr(activity_fn, LOCK_METADATA_KEY)
+            or IS_LOCKING_DISABLED
+        ):
+            return await self.next.start_activity(input)
+
+        lock_config = getattr(activity_fn, LOCK_METADATA_KEY)
+        lock_name = lock_config.get("lock_name", input.activity)
+        max_locks = lock_config.get("max_locks", 5)
+        if not input.schedule_to_close_timeout:
+            logger.error(
+                f"Activity '{input.activity}' with @needs_lock decorator requires schedule_to_close_timeout"
+            )
+            raise WorkflowError(
+                f"{WorkflowError.WORKFLOW_CONFIG_ERROR}: Activity '{input.activity}' with @needs_lock decorator must be called with schedule_to_close_timeout parameter. "
+                f"Example: workflow.execute_activity('{input.activity}', schedule_to_close_timeout=timedelta(minutes=10))"
+            )
+        ttl_seconds = int(input.schedule_to_close_timeout.total_seconds())
+
+        # Orchestrate lock acquisition -> business activity -> lock release
+        return await self._execute_with_lock_orchestration(
+            input, lock_name, max_locks, ttl_seconds
+        )
+
+    async def _execute_with_lock_orchestration(
+        self,
+        input: StartActivityInput,
+        lock_name: str,
+        max_locks: int,
+        ttl_seconds: int,
+    ) -> workflow.ActivityHandle[Any]:
+        """Execute activity with distributed lock orchestration."""
+        owner_id = f"{APPLICATION_NAME}:{workflow.info().run_id}"
+        lock_result = None
+
+        try:
+            # Step 1: Acquire lock via dedicated activity (can take >2s safely)
+            start_to_close_timeout = workflow.info().execution_timeout
+            lock_result = await workflow.execute_activity(
+                "acquire_distributed_lock",
+                args=[lock_name, max_locks, ttl_seconds, owner_id],
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=RetryPolicy(maximum_attempts=1),
+            )
+
+            logger.debug(f"Lock acquired: {lock_result}, executing {input.activity}")
+
+            # Step 2: Execute the business activity and return its handle
+            return await self.next.start_activity(input)
+
+        finally:
+            # Step 3: Release lock (fire-and-forget with short timeout)
+            if lock_result is not None:
+                try:
+                    await workflow.execute_local_activity(
+                        "release_distributed_lock",
+                        args=[lock_result["resource_id"], lock_result["owner_id"]],
+                        start_to_close_timeout=timedelta(seconds=5),
+                        retry_policy=RetryPolicy(maximum_attempts=1),
+                    )
+                    logger.debug(f"Lock released: {lock_result['resource_id']}")
+                except Exception as e:
+                    # Silent failure - TTL will handle cleanup
+                    logger.warning(
+                        f"Lock release failed for {lock_result['resource_id']}: {e}. "
+                        f"TTL will handle cleanup."
+                    )

application_sdk/outputs/__init__.py

@@ -13,6 +13,7 @@ from typing import (
     Dict,
     Generator,
     List,
+    Literal,
     Optional,
     Union,
     cast,
@@ -31,7 +32,7 @@ logger = get_logger(__name__)
 activity.logger = logger
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -52,6 +53,27 @@ class Output(ABC):
     output_prefix: str
     total_record_count: int
     chunk_count: int
+    statistics: List[int] = []
+
+    def estimate_dataframe_file_size(
+        self, dataframe: "pd.DataFrame", file_type: Literal["json", "parquet"]
+    ) -> 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)
+        if file_type == "json":
+            sample_file = sample.to_json(orient="records", lines=True)
+        else:
+            sample_file = sample.to_parquet(index=False, compression="snappy")
+        if sample_file is not None:
+            avg_record_size = len(sample_file) / sample_size
+            return int(avg_record_size * len(dataframe))
+
+        return 0
 
     def process_null_fields(
         self,
@@ -217,6 +239,7 @@ class Output(ABC):
             statistics = {
                 "total_record_count": self.total_record_count,
                 "chunk_count": self.chunk_count,
+                "partitions": self.statistics,
             }
 
             # Write the statistics to a json file

application_sdk/outputs/json.py

@@ -6,6 +6,7 @@ import orjson
 from temporalio import activity
 
 from application_sdk.activities.common.utils import get_object_store_prefix
+from application_sdk.constants import DAPR_MAX_GRPC_MESSAGE_LENGTH
 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
@@ -15,7 +16,7 @@ logger = get_logger(__name__)
 activity.logger = logger
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -32,7 +33,7 @@ def path_gen(chunk_start: int | None, chunk_count: int) -> str:
     if chunk_start is None:
         return f"{str(chunk_count)}.json"
     else:
-        return f"{str(chunk_start+chunk_count)}.json"
+        return f"chunk-{chunk_start}-part{chunk_count}.json"
 
 
 def convert_datetime_to_epoch(data: Any) -> Any:
@@ -124,6 +125,10 @@ class JsonOutput(Output):
         self.chunk_size = chunk_size or 100000
         self.buffer: List[Union["pd.DataFrame", "daft.DataFrame"]] = []  # noqa: F821
         self.current_buffer_size = 0
+        self.current_buffer_size_bytes = 0  # Track estimated buffer size in bytes
+        self.max_file_size_bytes = int(
+            DAPR_MAX_GRPC_MESSAGE_LENGTH * 0.9
+        )  # 90% of DAPR limit as safety buffer
         self.path_gen = path_gen
         self.start_marker = start_marker
         self.end_marker = end_marker
@@ -172,8 +177,21 @@ class JsonOutput(Output):
             ]
 
             for chunk in chunks:
+                # Estimate size of this chunk
+                chunk_size_bytes = self.estimate_dataframe_file_size(chunk, "json")
+
+                # Check if adding this chunk would exceed size limit
+                if (
+                    self.current_buffer_size_bytes + chunk_size_bytes
+                    > self.max_file_size_bytes
+                    and self.current_buffer_size > 0
+                ):
+                    # Flush current buffer before adding this chunk
+                    await self._flush_buffer()
+
                 self.buffer.append(chunk)
                 self.current_buffer_size += len(chunk)
+                self.current_buffer_size_bytes += chunk_size_bytes
 
                 if self.current_buffer_size >= partition:
                     await self._flush_buffer()
@@ -237,45 +255,19 @@ class JsonOutput(Output):
                     row, preserve_fields, null_to_empty_dict_fields
                 )
                 # Serialize the row and add it to the buffer
-                buffer.append(
-                    orjson.dumps(cleaned_row, option=orjson.OPT_APPEND_NEWLINE).decode(
-                        "utf-8"
-                    )
-                )
-
-            # If the buffer reaches the specified size, write it to the file
-            if self.chunk_size and len(buffer) >= self.chunk_size:
-                self.chunk_count += 1
-                output_file_name = f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count)}"
-                with open(output_file_name, "w") as f:
-                    f.writelines(buffer)
-                buffer.clear()  # Clear the buffer
-
-                # Record chunk metrics
-                self.metrics.record_metric(
-                    name="json_chunks_written",
-                    value=1,
-                    metric_type=MetricType.COUNTER,
-                    labels={"type": "daft"},
-                    description="Number of chunks written to JSON files",
-                )
+                serialized_row = orjson.dumps(
+                    cleaned_row, option=orjson.OPT_APPEND_NEWLINE
+                ).decode("utf-8")
+                buffer.append(serialized_row)
+                self.current_buffer_size_bytes += len(serialized_row)
+                if (self.chunk_size and len(buffer) >= self.chunk_size) or (
+                    self.current_buffer_size_bytes > self.max_file_size_bytes
+                ):
+                    await self.flush_daft_buffer(buffer)
 
             # Write any remaining rows in the buffer
             if buffer:
-                self.chunk_count += 1
-                output_file_name = f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count)}"
-                with open(output_file_name, "w") as f:
-                    f.writelines(buffer)
-                buffer.clear()
-
-                # Record chunk metrics
-                self.metrics.record_metric(
-                    name="json_chunks_written",
-                    value=1,
-                    metric_type=MetricType.COUNTER,
-                    labels={"type": "daft"},
-                    description="Number of chunks written to JSON files",
-                )
+                await self.flush_daft_buffer(buffer)
 
             # Record metrics for successful write
             self.metrics.record_metric(
@@ -303,6 +295,32 @@ class JsonOutput(Output):
             )
             logger.error(f"Error writing daft dataframe to json: {str(e)}")
 
+    async def flush_daft_buffer(self, buffer: List[str]):
+        """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.
+        """
+        self.chunk_count += 1
+        output_file_name = (
+            f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count)}"
+        )
+        with open(output_file_name, "w") as f:
+            f.writelines(buffer)
+        buffer.clear()  # Clear the buffer
+
+        self.current_buffer_size = 0
+        self.current_buffer_size_bytes = 0
+
+        # Record chunk metrics
+        self.metrics.record_metric(
+            name="json_chunks_written",
+            value=1,
+            metric_type=MetricType.COUNTER,
+            labels={"type": "daft"},
+            description="Number of chunks written to JSON files",
+        )
+
     async def _flush_buffer(self):
         """Flush the current buffer to a JSON file.
 
@@ -353,6 +371,7 @@ class JsonOutput(Output):
 
             self.buffer.clear()
             self.current_buffer_size = 0
+            self.current_buffer_size_bytes = 0
 
         except Exception as e:
             # Record metrics for failed write

application_sdk/outputs/parquet.py

@@ -1,9 +1,10 @@
 import os
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
 from temporalio import activity
 
 from application_sdk.activities.common.utils import get_object_store_prefix
+from application_sdk.constants import DAPR_MAX_GRPC_MESSAGE_LENGTH
 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
@@ -13,7 +14,7 @@ logger = get_logger(__name__)
 activity.logger = logger
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -46,6 +47,7 @@ class ParquetOutput(Output):
         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,
         chunk_count: int = 0,
         chunk_start: Optional[int] = None,
@@ -78,11 +80,19 @@ class ParquetOutput(Output):
         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
         self.total_record_count = total_record_count
         self.chunk_count = chunk_count
+        self.current_buffer_size = 0
+        self.current_buffer_size_bytes = 0  # Track estimated buffer size in bytes
+        self.max_file_size_bytes = int(
+            DAPR_MAX_GRPC_MESSAGE_LENGTH * 0.9
+        )  # 90% of DAPR limit as safety buffer
         self.chunk_start = chunk_start
         self.start_marker = start_marker
         self.end_marker = end_marker
+        self.statistics = []
         self.metrics = get_metrics()
 
         # Create output directory
@@ -117,7 +127,7 @@ class ParquetOutput(Output):
         if chunk_start is None:
             return f"{str(chunk_count)}.parquet"
         else:
-            return f"{str(chunk_start+chunk_count)}.parquet"
+            return f"chunk-{str(chunk_start)}-part{str(chunk_count)}.parquet"
 
     async def write_dataframe(self, dataframe: "pd.DataFrame"):
         """Write a pandas DataFrame to Parquet files and upload to object store.
@@ -126,20 +136,46 @@ class ParquetOutput(Output):
             dataframe (pd.DataFrame): The DataFrame to write.
         """
         try:
+            chunk_part = 0
             if len(dataframe) == 0:
                 return
 
-            # Update counters
-            self.chunk_count += 1
-            self.total_record_count += len(dataframe)
-            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 pandas native method
-            dataframe.to_parquet(
-                file_path,
-                index=False,
-                compression="snappy",  # Using snappy compression by default
+            # Split the DataFrame into chunks
+            partition = (
+                self.chunk_size
+                if self.chunk_start is None
+                else min(self.chunk_size, self.buffer_size)
             )
+            chunks = [
+                dataframe[i : i + partition]  # type: ignore
+                for i in range(0, len(dataframe), partition)
+            ]
+
+            for chunk in chunks:
+                # Estimate size of this chunk
+                chunk_size_bytes = self.estimate_dataframe_file_size(chunk, "parquet")
+
+                # Check if adding this chunk would exceed size limit
+                if (
+                    self.current_buffer_size_bytes + chunk_size_bytes
+                    > self.max_file_size_bytes
+                    and self.current_buffer_size > 0
+                ):
+                    # Flush current buffer before adding this chunk
+                    chunk_part += 1
+                    await self._flush_buffer(chunk_part)
+
+                self.buffer.append(chunk)
+                self.current_buffer_size += len(chunk)
+                self.current_buffer_size_bytes += chunk_size_bytes
+
+                if self.current_buffer_size >= partition:  # type: ignore
+                    chunk_part += 1
+                    await self._flush_buffer(chunk_part)
+
+            if self.buffer and self.current_buffer_size > 0:
+                chunk_part += 1
+                await self._flush_buffer(chunk_part)
 
             # Record metrics for successful write
             self.metrics.record_metric(
@@ -159,11 +195,8 @@ class ParquetOutput(Output):
                 description="Number of chunks written to Parquet files",
             )
 
-            # Upload the file to object store
-            await ObjectStore.upload_file(
-                source=file_path,
-                destination=get_object_store_prefix(file_path),
-            )
+            self.chunk_count += 1
+            self.statistics.append(chunk_part)
         except Exception as e:
             # Record metrics for failed write
             self.metrics.record_metric(
@@ -245,3 +278,68 @@ class ParquetOutput(Output):
             str: The full path of the output file.
         """
         return self.output_path
+
+    async def _flush_buffer(self, chunk_part):
+        """Flush the current buffer to a Parquet file.
+
+        This method combines all DataFrames in the buffer, writes them to a Parquet file,
+        and uploads the file to the object store.
+
+        Note:
+            If the buffer is empty or has no records, the method returns without writing.
+        """
+        import pandas as pd
+
+        if not self.buffer or not self.current_buffer_size:
+            return
+
+        if not all(isinstance(df, pd.DataFrame) for df in self.buffer):
+            raise TypeError(
+                "_flush_buffer encountered non-DataFrame elements in buffer. This should not happen."
+            )
+
+        try:
+            # Now it's safe to cast for pd.concat
+            pd_buffer: List[pd.DataFrame] = self.buffer  # type: ignore
+            combined_dataframe = pd.concat(pd_buffer)
+
+            # Write DataFrame to Parquet file
+            if not combined_dataframe.empty:
+                self.total_record_count += len(combined_dataframe)
+                output_file_name = (
+                    f"{self.output_path}/{self.path_gen(self.chunk_count, chunk_part)}"
+                )
+                combined_dataframe.to_parquet(
+                    output_file_name, index=False, compression="snappy"
+                )
+
+                # Record chunk metrics
+                self.metrics.record_metric(
+                    name="parquet_chunks_written",
+                    value=1,
+                    metric_type=MetricType.COUNTER,
+                    labels={"type": "pandas"},
+                    description="Number of chunks written to Parquet files",
+                )
+
+                # Push the file to the object store
+                await ObjectStore.upload_file(
+                    source=output_file_name,
+                    destination=get_object_store_prefix(output_file_name),
+                )
+
+            self.buffer.clear()
+            self.current_buffer_size = 0
+            self.current_buffer_size_bytes = 0
+
+        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": "pandas", "error": str(e)},
+                description="Number of errors while writing to Parquet files",
+            )
+            logger.error(f"Error flushing buffer to parquet: {str(e)}")
+            raise e

application_sdk/version.py

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

application_sdk/workflows/metadata_extraction/sql.py

@@ -107,7 +107,11 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
         activity_statistics = ActivityStatistics.model_validate(raw_statistics)
         transform_activities: List[Any] = []
 
-        if activity_statistics is None or activity_statistics.chunk_count == 0:
+        if (
+            activity_statistics is None
+            or activity_statistics.chunk_count == 0
+            or not activity_statistics.partitions
+        ):
             # to handle the case where the fetch_fn returns None or no chunks
             return
 
@@ -115,7 +119,9 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
             raise ValueError("Invalid typename")
 
         batches, chunk_starts = self.get_transform_batches(
-            activity_statistics.chunk_count, activity_statistics.typename
+            activity_statistics.chunk_count,
+            activity_statistics.typename,
+            activity_statistics.partitions,
         )
 
         for i in range(len(batches)):
@@ -144,7 +150,9 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
             total_record_count += metadata_model.total_record_count
             chunk_count += metadata_model.chunk_count
 
-    def get_transform_batches(self, chunk_count: int, typename: str):
+    def get_transform_batches(
+        self, chunk_count: int, typename: str, partitions: List[int]
+    ):
         """Get batches for parallel transformation processing.
 
         Args:
@@ -159,12 +167,17 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
         batches: List[List[str]] = []
         chunk_start_numbers: List[int] = []
 
-        for i in range(chunk_count):
+        for i, partition in enumerate(partitions):
             # Track starting chunk number (which is just i)
             chunk_start_numbers.append(i)
 
             # Each batch contains exactly one chunk
-            batches.append([f"{typename}/{i+1}.json"])
+            batches.append(
+                [
+                    f"{typename}/chunk-{i}-part{file+1}.parquet"
+                    for file in range(partition)
+                ]
+            )
 
         return batches, chunk_start_numbers
 

{atlan_application_sdk-0.1.1rc35.dist-info → atlan_application_sdk-0.1.1rc37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc35
+Version: 0.1.1rc37
 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
@@ -26,12 +26,14 @@ Requires-Dist: fastapi[standard]>=0.115.0
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: opentelemetry-exporter-otlp>=1.27.0
 Requires-Dist: psutil>=7.0.0
-Requires-Dist: pyatlan>=8.0.0
+Requires-Dist: pyatlan>=8.0.2
 Requires-Dist: pydantic>=2.10.6
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: uvloop>=0.21.0; sys_platform != 'win32'
 Provides-Extra: daft
 Requires-Dist: daft>=0.4.12; extra == 'daft'
+Provides-Extra: distributed-lock
+Requires-Dist: redis[hiredis]>=5.2.0; extra == 'distributed-lock'
 Provides-Extra: iam-auth
 Requires-Dist: boto3>=1.38.6; extra == 'iam-auth'
 Provides-Extra: iceberg