atlan-application-sdk 0.1.1rc43__py3-none-any.whl → 0.1.1rc45__py3-none-any.whl

application_sdk/activities/metadata_extraction/sql.py

@@ -2,10 +2,8 @@ import os
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncGenerator,
     AsyncIterator,
     Dict,
-    Generator,
     Iterator,
     List,
     Optional,
@@ -368,9 +366,9 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
                 "Output prefix and path must be specified in workflow_args."
             )
         return ParquetOutput(
-            output_prefix=output_prefix,
             output_path=output_path,
             output_suffix=output_suffix,
+            use_consolidation=True,
         )
 
     def _get_temp_table_regex_sql(self, typename: str) -> str:
@@ -553,7 +551,7 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
                 )
 
                 # Execute using helper method
-                success, batched_iter = await self._execute_single_db(
+                success, batched_iterator = await self._execute_single_db(
                     effective_sql_client.engine,
                     prepared_query,
                     parquet_output,
@@ -570,12 +568,12 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
                 logger.warning(
                     f"Failed to process database '{database_name}': {str(e)}. Skipping to next database."
                 )
-                success, batched_iter = False, None
+                success, batched_iterator = False, None
 
             if success:
                 successful_databases.append(database_name)
-                if not write_to_file and batched_iter:
-                    dataframe_list.append(batched_iter)
+                if not write_to_file and batched_iterator:
+                    dataframe_list.append(batched_iterator)
             else:
                 failed_databases.append(database_name)
 
@@ -615,37 +613,13 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
 
         try:
             sql_input = SQLQueryInput(engine=sql_engine, query=prepared_query)
-            batched_iter = await sql_input.get_batched_dataframe()
+            batched_iterator = await sql_input.get_batched_dataframe()
 
             if write_to_file and parquet_output:
-                # Wrap iterator into a proper (async)generator for type safety
-                if hasattr(batched_iter, "__anext__"):
-
-                    async def _to_async_gen(
-                        it: AsyncIterator["pd.DataFrame"],
-                    ) -> AsyncGenerator["pd.DataFrame", None]:
-                        async for item in it:
-                            yield item
-
-                    wrapped: AsyncGenerator["pd.DataFrame", None] = _to_async_gen(  # type: ignore
-                        batched_iter  # type: ignore
-                    )
-                    await parquet_output.write_batched_dataframe(wrapped)
-                else:
-
-                    def _to_gen(
-                        it: Iterator["pd.DataFrame"],
-                    ) -> Generator["pd.DataFrame", None, None]:
-                        for item in it:
-                            yield item
-
-                    wrapped_sync: Generator["pd.DataFrame", None, None] = _to_gen(  # type: ignore
-                        batched_iter  # type: ignore
-                    )
-                    await parquet_output.write_batched_dataframe(wrapped_sync)
+                await parquet_output.write_batched_dataframe(batched_iterator)  # type: ignore
                 return True, None
 
-            return True, batched_iter
+            return True, batched_iterator
         except Exception as e:
             logger.error(
                 f"Error during query execution or output writing: {e}", exc_info=True
@@ -863,10 +837,10 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
             file_names=workflow_args.get("file_names"),
         )
         raw_input = raw_input.get_batched_daft_dataframe()
+
         transformed_output = JsonOutput(
             output_path=output_path,
             output_suffix="transformed",
-            output_prefix=output_prefix,
             typename=typename,
             chunk_start=workflow_args.get("chunk_start"),
         )

application_sdk/activities/query_extraction/sql.py

@@ -210,7 +210,6 @@ class SQLQueryExtractionActivities(ActivitiesInterface):
             sql_input = await sql_input.get_dataframe()
 
             raw_output = ParquetOutput(
-                output_prefix=workflow_args["output_prefix"],
                 output_path=workflow_args["output_path"],
                 output_suffix="raw/query",
                 chunk_size=workflow_args["miner_args"].get("chunk_size", 100000),
@@ -218,7 +217,6 @@ class SQLQueryExtractionActivities(ActivitiesInterface):
                 end_marker=workflow_args["end_marker"],
             )
             await raw_output.write_dataframe(sql_input)
-
             logger.info(
                 f"Query fetch completed, {raw_output.total_record_count} records processed",
             )

application_sdk/application/__init__.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, Type
 from application_sdk.activities import ActivitiesInterface
 from application_sdk.clients.base import BaseClient
 from application_sdk.clients.utils import get_workflow_client
+from application_sdk.constants import ENABLE_MCP
 from application_sdk.events.models import EventRegistration
 from application_sdk.handlers.base import BaseHandler
 from application_sdk.observability.logger_adaptor import get_logger
@@ -29,7 +30,7 @@ class BaseApplication:
         self,
         name: str,
         server: Optional[ServerInterface] = None,
-        application_manifest: Optional[dict] = None,
+        application_manifest: Optional[Dict[str, Any]] = None,
         client_class: Optional[Type[BaseClient]] = None,
         handler_class: Optional[Type[BaseHandler]] = None,
     ):
@@ -39,6 +40,9 @@ class BaseApplication:
         Args:
             name (str): The name of the application.
             server (ServerInterface): The server class for the application.
+            application_manifest (Optional[Dict[str, Any]]): Application manifest configuration.
+            client_class (Optional[Type[BaseClient]]): Client class for the application.
+            handler_class (Optional[Type[BaseHandler]]): Handler class for the application.
         """
         self.application_name = name
 
@@ -49,14 +53,21 @@ class BaseApplication:
 
         self.workflow_client = get_workflow_client(application_name=name)
 
-        self.application_manifest: Dict[str, Any] = application_manifest
+        self.application_manifest: Optional[Dict[str, Any]] = application_manifest
         self.bootstrap_event_registration()
 
         self.client_class = client_class or BaseClient
         self.handler_class = handler_class or BaseHandler
 
+        # MCP configuration
+        self.mcp_server: Optional["MCPServer"] = None
+        if ENABLE_MCP:
+            from application_sdk.server.mcp import MCPServer
+
+            self.mcp_server = MCPServer(application_name=name)
+
     def bootstrap_event_registration(self):
-        self.event_subscriptions = {}
+        self.event_subscriptions: Dict[str, EventWorkflowTrigger] = {}
         if self.application_manifest is None:
             logger.warning("No application manifest found, skipping event registration")
             return
@@ -122,8 +133,8 @@ class BaseApplication:
         ]
         workflow_activities = []
         for workflow_class, activities_class in workflow_and_activities_classes:
-            workflow_activities.extend(
-                workflow_class.get_activities(activities_class())
+            workflow_activities.extend(  # type: ignore
+                workflow_class.get_activities(activities_class())  # type: ignore
             )
 
         self.worker = Worker(
@@ -134,6 +145,13 @@ class BaseApplication:
             activity_executor=activity_executor,
         )
 
+        # Register MCP tools if ENABLED_MCP is True and an MCP server is initialized
+        if self.mcp_server:
+            logger.info("Registering MCP tools from workflow and activities classes")
+            await self.mcp_server.register_tools(  # type: ignore
+                workflow_and_activities_classes=workflow_and_activities_classes
+            )
+
     async def start_workflow(self, workflow_args, workflow_class) -> Any:
         """
         Start a new workflow execution.
@@ -147,7 +165,7 @@ class BaseApplication:
         """
         if self.workflow_client is None:
             raise ValueError("Workflow client not initialized")
-        return await self.workflow_client.start_workflow(workflow_args, workflow_class)
+        return await self.workflow_client.start_workflow(workflow_args, workflow_class)  # type: ignore
 
     async def start_worker(self, daemon: bool = True):
         """
@@ -162,39 +180,61 @@ class BaseApplication:
 
     async def setup_server(
         self,
-        workflow_class,
+        workflow_class: Type[WorkflowInterface],
         ui_enabled: bool = True,
         has_configmap: bool = False,
     ):
         """
-        Optionally set up a server for the application. (No-op by default)
+        Set up FastAPI server and automatically mount MCP if enabled.
+
+        Args:
+            workflow_class (WorkflowInterface): The workflow class for the application.
+            ui_enabled (bool): Whether to enable the UI.
+            has_configmap (bool): Whether to enable the configmap.
         """
         if self.workflow_client is None:
             await self.workflow_client.load()
 
-        # Overrides the application server. serves the UI, and handles the various triggers
+        mcp_http_app: Optional[Any] = None
+        lifespan: Optional[Any] = None
+
+        if self.mcp_server:
+            try:
+                mcp_http_app = await self.mcp_server.get_http_app()
+                lifespan = mcp_http_app.lifespan
+            except Exception as e:
+                logger.warning(f"Failed to get MCP HTTP app: {e}")
+
         self.server = APIServer(
+            lifespan=lifespan,
             workflow_client=self.workflow_client,
             ui_enabled=ui_enabled,
             handler=self.handler_class(client=self.client_class()),
             has_configmap=has_configmap,
         )
 
+        # Mount MCP at root
+        if mcp_http_app:
+            try:
+                self.server.app.mount("", mcp_http_app)  # Mount at root
+            except Exception as e:
+                logger.warning(f"Failed to mount MCP HTTP app: {e}")
+
+        # Register event-based workflows if any
         if self.event_subscriptions:
             for event_trigger in self.event_subscriptions.values():
-                if event_trigger.workflow_class is None:
+                if event_trigger.workflow_class is None:  # type: ignore
                     raise ValueError(
                         f"Workflow class not set for event trigger {event_trigger.event_id}"
                     )
 
-                self.server.register_workflow(
-                    workflow_class=event_trigger.workflow_class,
+                self.server.register_workflow(  # type: ignore
+                    workflow_class=event_trigger.workflow_class,  # type: ignore
                     triggers=[event_trigger],
                 )
 
-        # register the workflow on the application server
-        # the workflow is by default triggered by an HTTP POST request to the /start endpoint
-        self.server.register_workflow(
+        # Register the main workflow (HTTP POST /start endpoint)
+        self.server.register_workflow(  # type: ignore
             workflow_class=workflow_class,
             triggers=[HttpWorkflowTrigger()],
         )

application_sdk/constants.py

@@ -255,3 +255,9 @@ REDIS_SENTINEL_HOSTS = os.getenv("REDIS_SENTINEL_HOSTS", "")
 IS_LOCKING_DISABLED = os.getenv("IS_LOCKING_DISABLED", "true").lower() == "true"
 #: Retry interval for lock acquisition
 LOCK_RETRY_INTERVAL = int(os.getenv("LOCK_RETRY_INTERVAL", "5"))
+
+# MCP Configuration
+#: Flag to indicate if MCP should be enabled or not. Turning this to true will setup an MCP server along
+#: with the application.
+ENABLE_MCP = os.getenv("ENABLE_MCP", "false").lower() == "true"
+MCP_METADATA_KEY = "__atlan_application_sdk_mcp_metadata"

application_sdk/decorators/mcp_tool.py

@@ -0,0 +1,63 @@
+"""
+MCP tool decorator for marking activities as MCP tools.
+
+This module provides the @mcp_tool decorator that developers use to mark
+activities for automatic exposure via Model Context Protocol.
+"""
+
+from typing import Any, Callable, Optional
+
+from application_sdk.constants import MCP_METADATA_KEY
+from application_sdk.server.mcp import MCPMetadata
+
+
+def mcp_tool(
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    visible: bool = True,
+    *args,
+    **kwargs,
+):
+    """
+    Decorator to mark functions as MCP tools.
+
+    Use this decorator to mark any function as an MCP tool. You can additionally use the `visible`
+    parameter to control whether the tool is visible at runtime or not.
+
+    Function parameters that are Pydantic models will be automatically converted into correct JSON schema
+    for the tool specification. This is handled by the underlying FastMCP server implementation.
+
+    Args:
+        name(Optional[str]): The name of the tool. Defaults to the function name.
+        description(Optional[str]): The description of the tool. Defaults to the function docstring.
+        visible(bool): Whether the MCP tool is visible at runtime or not. Defaults to True.
+        *args: Additional arguments to pass to the tool.
+        **kwargs: Additional keyword arguments to pass to the tool.
+
+    Examples:
+        >>> @mcp_tool(name="add_numbers", description="Add two numbers", visible=True)
+        >>> def add_numbers(self, a: int, b: int) -> int:
+        >>>     return a + b
+
+
+        >>> # Use with Temporal activity decorator
+        >>> @activity.defn
+        >>> @mcp_tool(name="get_weather", description="Get the weather for a given city")
+        >>> async def get_weather(self, city: str) -> str:
+        >>>     # ... activity implementation unchanged ...
+    """
+
+    def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        mcp_metadata = MCPMetadata(
+            name=name if name else f.__name__,
+            description=description if description else f.__doc__,
+            visible=visible,
+            args=args,
+            kwargs=kwargs,
+        )
+
+        setattr(f, MCP_METADATA_KEY, mcp_metadata)
+
+        return f
+
+    return decorator

application_sdk/inputs/parquet.py

@@ -22,6 +22,7 @@ class ParquetInput(Input):
         self,
         path: str,
         chunk_size: int = 100000,
+        buffer_size: int = 5000,
         file_names: Optional[List[str]] = None,
     ):
         """Initialize the Parquet input class.
@@ -32,6 +33,7 @@ class ParquetInput(Input):
                 local path or object store path
                 Wildcards are not supported.
             chunk_size (int): Number of rows per batch. Defaults to 100000.
+            buffer_size (int): Number of rows per batch. Defaults to 5000.
             file_names (Optional[List[str]]): List of file names to read. Defaults to None.
 
         Raises:
@@ -47,6 +49,7 @@ class ParquetInput(Input):
 
         self.path = path
         self.chunk_size = chunk_size
+        self.buffer_size = buffer_size
         self.file_names = file_names
 
     async def get_dataframe(self) -> "pd.DataFrame":
@@ -249,9 +252,18 @@ class ParquetInput(Input):
             parquet_files = await self.download_files()
             logger.info(f"Reading {len(parquet_files)} parquet files as daft batches")
 
-            # Yield each discovered file as separate batch
-            for parquet_file in parquet_files:
-                yield daft.read_parquet(parquet_file)
+            # Create a lazy dataframe without loading data into memory
+            lazy_df = daft.read_parquet(parquet_files)
+
+            # Get total count efficiently
+            total_rows = lazy_df.count_rows()
+
+            # Yield chunks without loading everything into memory
+            for offset in range(0, total_rows, self.buffer_size):
+                chunk = lazy_df.offset(offset).limit(self.buffer_size)
+                yield chunk
+
+            del lazy_df
 
         except Exception as error:
             logger.error(

application_sdk/inputs/sql_query.py

@@ -34,7 +34,7 @@ class SQLQueryInput(Input):
         self,
         query: str,
         engine: Union["Engine", str],
-        chunk_size: Optional[int] = 100000,
+        chunk_size: Optional[int] = 5000,
     ):
         """Initialize the async SQL query input handler.
 
@@ -42,7 +42,7 @@ class SQLQueryInput(Input):
             engine (Union[Engine, str]): SQLAlchemy engine or connection string.
             query (str): The SQL query to execute.
             chunk_size (Optional[int], optional): Number of rows per batch.
-                Defaults to 100000.
+                Defaults to 5000.
         """
         self.query = query
         self.engine = engine

application_sdk/interceptors/cleanup.py

@@ -114,7 +114,6 @@ class CleanupWorkflowInboundInterceptor(WorkflowInboundInterceptor):
                     retry_policy=RetryPolicy(
                         maximum_attempts=3,
                     ),
-                    summary="This activity is used to cleanup the local artifacts and the activity state after the workflow is completed.",
                 )
 
                 logger.info("Cleanup completed successfully")

application_sdk/outputs/__init__.py

@@ -4,8 +4,11 @@ This module provides base classes and utilities for handling various types of da
 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,
@@ -13,7 +16,6 @@ from typing import (
     Dict,
     Generator,
     List,
-    Literal,
     Optional,
     Union,
     cast,
@@ -26,6 +28,7 @@ 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.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__)
@@ -36,6 +39,14 @@ if TYPE_CHECKING:
     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.
 
@@ -53,11 +64,13 @@ class Output(ABC):
     output_prefix: str
     total_record_count: int
     chunk_count: int
-    statistics: List[int] = []
+    buffer_size: int
+    max_file_size_bytes: int
+    current_buffer_size: int
+    current_buffer_size_bytes: int
+    partitions: List[int]
 
-    def estimate_dataframe_file_size(
-        self, dataframe: "pd.DataFrame", file_type: Literal["json", "parquet"]
-    ) -> 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
@@ -65,16 +78,47 @@ class Output(ABC):
         # 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
-            return int(avg_record_size * len(dataframe))
+            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,
@@ -146,15 +190,86 @@ class Output(ABC):
                         await self.write_dataframe(dataframe)
         except Exception as e:
             logger.error(f"Error writing batched dataframe: {str(e)}")
+            raise
 
-    @abstractmethod
     async def write_dataframe(self, dataframe: "pd.DataFrame"):
-        """Write a pandas DataFrame to the output destination.
+        """Write a pandas DataFrame to Parquet files and upload to object store.
 
         Args:
             dataframe (pd.DataFrame): The DataFrame to write.
         """
-        pass
+        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,
@@ -225,6 +340,55 @@ class Output(ABC):
             logger.error(f"Error getting statistics: {str(e)}")
             raise
 
+    async def _upload_file(self, file_name: str):
+        """Upload a file to the object store."""
+        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) -> Optional[Dict[str, Any]]:
         """Write statistics about the output to a JSON file.
 
@@ -238,8 +402,8 @@ class Output(ABC):
             # prepare the statistics
             statistics = {
                 "total_record_count": self.total_record_count,
-                "chunk_count": self.chunk_count,
-                "partitions": self.statistics,
+                "chunk_count": len(self.partitions),
+                "partitions": self.partitions,
             }
 
             # Write the statistics to a json file