unstructured-ingest 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

unstructured_ingest/v2/interfaces/upload_stager.py

@@ -1,8 +1,10 @@
-from abc import ABC, abstractmethod
+import json
+from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, TypeVar
 
+import ndjson
 from pydantic import BaseModel
 
 from unstructured_ingest.v2.interfaces.file_data import FileData
@@ -20,16 +22,78 @@ UploadStagerConfigT = TypeVar("UploadStagerConfigT", bound=UploadStagerConfig)
 class UploadStager(BaseProcess, ABC):
     upload_stager_config: UploadStagerConfigT
 
-    @abstractmethod
+    def write_output(self, output_path: Path, data: list[dict]) -> None:
+        if output_path.suffix == ".json":
+            with output_path.open("w") as f:
+                json.dump(data, f, indent=2)
+        elif output_path.suffix == ".ndjson":
+            with output_path.open("w") as f:
+                ndjson.dump(data, f)
+        else:
+            raise ValueError(f"Unsupported output format: {output_path}")
+
+    def get_data(self, elements_filepath: Path) -> list[dict]:
+        if elements_filepath.suffix == ".json":
+            with elements_filepath.open() as f:
+                return json.load(f)
+        elif elements_filepath.suffix == ".ndjson":
+            with elements_filepath.open() as f:
+                return ndjson.load(f)
+        else:
+            raise ValueError(f"Unsupported input format: {elements_filepath}")
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        return element_dict
+
+    def get_output_path(self, output_filename: str, output_dir: Path) -> Path:
+        output_path = Path(output_filename)
+        output_filename = f"{Path(output_filename).stem}{output_path.suffix}"
+        output_path = Path(output_dir) / Path(f"{output_filename}")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        return output_path
+
+    def stream_update(self, input_file: Path, output_file: Path, file_data: FileData) -> None:
+        with input_file.open() as in_f:
+            reader = ndjson.reader(in_f)
+            with output_file.open("w") as out_f:
+                writer = ndjson.writer(out_f)
+                for element in reader:
+                    conformed_element = self.conform_dict(element_dict=element, file_data=file_data)
+                    writer.writerow(row=conformed_element)
+                    writer.f.flush()
+
+    def process_whole(self, input_file: Path, output_file: Path, file_data: FileData) -> None:
+        with input_file.open() as in_f:
+            elements_contents = json.load(in_f)
+
+        conformed_elements = [
+            self.conform_dict(element_dict=element, file_data=file_data)
+            for element in elements_contents
+        ]
+
+        with open(output_file, "w") as out_f:
+            json.dump(conformed_elements, out_f, indent=2)
+
     def run(
         self,
         elements_filepath: Path,
         file_data: FileData,
         output_dir: Path,
         output_filename: str,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> Path:
-        pass
+        output_file = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
+        if elements_filepath.suffix == ".ndjson":
+            self.stream_update(
+                input_file=elements_filepath, output_file=output_file, file_data=file_data
+            )
+        elif elements_filepath.suffix == ".json":
+            self.process_whole(
+                input_file=elements_filepath, output_file=output_file, file_data=file_data
+            )
+        else:
+            raise ValueError(f"Unsupported file extension: {elements_filepath}")
+        return output_file
 
     async def run_async(
         self,
@@ -37,12 +101,12 @@ class UploadStager(BaseProcess, ABC):
         file_data: FileData,
         output_dir: Path,
         output_filename: str,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> Path:
         return self.run(
             elements_filepath=elements_filepath,
             output_dir=output_dir,
             output_filename=output_filename,
             file_data=file_data,
-            **kwargs
+            **kwargs,
         )

unstructured_ingest/v2/interfaces/uploader.py

@@ -5,6 +5,7 @@ from typing import Any, TypeVar
 
 from pydantic import BaseModel
 
+from unstructured_ingest.utils.data_prep import get_data
 from unstructured_ingest.v2.interfaces.connector import BaseConnector
 from unstructured_ingest.v2.interfaces.file_data import FileData
 from unstructured_ingest.v2.interfaces.process import BaseProcess
@@ -38,7 +39,15 @@ class Uploader(BaseProcess, BaseConnector, ABC):
         raise NotImplementedError()
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        raise NotImplementedError()
+        data = get_data(path=path)
+        self.run_data(data=data, file_data=file_data, **kwargs)
 
     async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return self.run(contents=[UploadContent(path=path, file_data=file_data)], **kwargs)
+        data = get_data(path=path)
+        await self.run_data_async(data=data, file_data=file_data, **kwargs)
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        raise NotImplementedError()
+
+    async def run_data_async(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        return self.run_data(data=data, file_data=file_data, **kwargs)

unstructured_ingest/v2/pipeline/steps/chunk.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.chunker import Chunker
@@ -51,7 +52,7 @@ class ChunkStep(PipelineStep):
         self, fn: Callable, path: str, file_data_path: str, **kwargs
     ) -> ChunkStepResponse:
         path = Path(path)
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         output_filepath = self.get_output_filepath(filename=path)
         if not self.should_chunk(filepath=output_filepath, file_data=file_data):
             logger.debug(f"skipping chunking, output already exists: {output_filepath}")

unstructured_ingest/v2/pipeline/steps/download.py

@@ -8,6 +8,7 @@ from typing import Callable, Optional, TypedDict, TypeVar
 
 from unstructured_ingest.v2.interfaces import FileData, download_responses
 from unstructured_ingest.v2.interfaces.downloader import Downloader
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.utils import serialize_base_model_json
@@ -87,12 +88,12 @@ class DownloadStep(PipelineStep):
                 f"match size of local file: {file_size_bytes}, updating"
             )
             file_data.metadata.filesize_bytes = file_size_bytes
-        logger.debug(f"updating file data with new content: {file_data.to_dict()}")
+        logger.debug(f"updating file data with new content: {file_data.model_dump()}")
         with file_data_path.open("w") as file:
-            json.dump(file_data.to_dict(), file, indent=2)
+            json.dump(file_data.model_dump(), file, indent=2)
 
     async def _run_async(self, fn: Callable, file_data_path: str) -> list[DownloadStepResponse]:
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         download_path = self.process.get_download_path(file_data=file_data)
         if not self.should_download(file_data=file_data, file_data_path=file_data_path):
             logger.debug(f"skipping download, file already exists locally: {download_path}")
@@ -172,7 +173,7 @@ class DownloadStep(PipelineStep):
         filepath = (self.cache_dir / filename).resolve()
         filepath.parent.mkdir(parents=True, exist_ok=True)
         with open(str(filepath), "w") as f:
-            json.dump(file_data.to_dict(), f, indent=2)
+            json.dump(file_data.model_dump(), f, indent=2)
         return str(filepath)
 
     def get_hash(self, extras: Optional[list[str]]) -> str:

unstructured_ingest/v2/pipeline/steps/embed.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.embedder import Embedder
@@ -49,7 +50,7 @@ class EmbedStep(PipelineStep):
 
     async def _run_async(self, fn: Callable, path: str, file_data_path: str) -> EmbedStepResponse:
         path = Path(path)
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         output_filepath = self.get_output_filepath(filename=path)
         if not self.should_embed(filepath=output_filepath, file_data=file_data):
             logger.debug(f"skipping embedding, output already exists: {output_filepath}")

unstructured_ingest/v2/pipeline/steps/filter.py

@@ -2,7 +2,7 @@ import asyncio
 from dataclasses import dataclass
 from typing import Callable, Optional
 
-from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.filter import Filterer
@@ -20,7 +20,7 @@ class FilterStep(PipelineStep):
         logger.info(f"created {self.identifier} with configs: {config}")
 
     async def _run_async(self, fn: Callable, file_data_path: str, **kwargs) -> Optional[dict]:
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         fn_kwargs = {"file_data": file_data}
         if not asyncio.iscoroutinefunction(fn):
             resp = fn(**fn_kwargs)

unstructured_ingest/v2/pipeline/steps/index.py

@@ -37,14 +37,14 @@ class IndexStep(PipelineStep):
     @instrument(span_name=STEP_ID)
     def run(self) -> Generator[str, None, None]:
         for file_data in self.process.run():
-            logger.debug(f"generated file data: {file_data.to_dict()}")
+            logger.debug(f"generated file data: {file_data.model_dump()}")
             try:
                 record_hash = self.get_hash(extras=[file_data.identifier])
                 filename = f"{record_hash}.json"
                 filepath = (self.cache_dir / filename).resolve()
                 filepath.parent.mkdir(parents=True, exist_ok=True)
                 with open(str(filepath), "w") as f:
-                    json.dump(file_data.to_dict(), f, indent=2)
+                    json.dump(file_data.model_dump(), f, indent=2)
                 yield str(filepath)
             except Exception as e:
                 logger.error(f"failed to create index for file data: {file_data}", exc_info=True)
@@ -54,14 +54,14 @@ class IndexStep(PipelineStep):
 
     async def run_async(self) -> AsyncGenerator[str, None]:
         async for file_data in self.process.run_async():
-            logger.debug(f"generated file data: {file_data.to_dict()}")
+            logger.debug(f"generated file data: {file_data.model_dump()}")
             try:
                 record_hash = self.get_hash(extras=[file_data.identifier])
                 filename = f"{record_hash}.json"
                 filepath = (self.cache_dir / filename).resolve()
                 filepath.parent.mkdir(parents=True, exist_ok=True)
                 with open(str(filepath), "w") as f:
-                    json.dump(file_data.to_dict(), f, indent=2)
+                    json.dump(file_data.model_dump(), f, indent=2)
                 yield str(filepath)
             except Exception as e:
                 logger.error(f"failed to create index for file data: {file_data}", exc_info=True)

unstructured_ingest/v2/pipeline/steps/partition.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.partitioner import Partitioner
@@ -51,12 +52,12 @@ class PartitionStep(PipelineStep):
         self, fn: Callable, path: str, file_data_path: str
     ) -> Optional[PartitionStepResponse]:
         path = Path(path)
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         output_filepath = self.get_output_filepath(filename=Path(file_data_path))
         if not self.should_partition(filepath=output_filepath, file_data=file_data):
             logger.debug(f"skipping partitioning, output already exists: {output_filepath}")
             return PartitionStepResponse(file_data_path=file_data_path, path=str(output_filepath))
-        fn_kwargs = {"filename": path, "metadata": file_data.metadata.to_dict()}
+        fn_kwargs = {"filename": path, "metadata": file_data.metadata.model_dump()}
         if not asyncio.iscoroutinefunction(fn):
             partitioned_content = fn(**fn_kwargs)
         elif semaphore := self.context.semaphore:

unstructured_ingest/v2/pipeline/steps/stage.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
-from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.interfaces.upload_stager import UploadStager
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
@@ -39,11 +39,13 @@ class UploadStageStep(PipelineStep):
         self, fn: Callable, path: str, file_data_path: str
     ) -> UploadStageStepResponse:
         path = Path(path)
+        # Maintain extension
+        output_filename = f"{self.get_hash(extras=[path.name])}{path.suffix}"
         fn_kwargs = {
             "elements_filepath": path,
-            "file_data": FileData.from_file(path=file_data_path),
+            "file_data": file_data_from_file(path=file_data_path),
             "output_dir": self.cache_dir,
-            "output_filename": self.get_hash(extras=[path.name]),
+            "output_filename": output_filename,
         }
         if not asyncio.iscoroutinefunction(fn):
             staged_output_path = fn(**fn_kwargs)

unstructured_ingest/v2/pipeline/steps/uncompress.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, TypedDict
 
-from unstructured_ingest.v2.interfaces.file_data import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import PipelineStep
 from unstructured_ingest.v2.processes.uncompress import Uncompressor
@@ -28,7 +28,7 @@ class UncompressStep(PipelineStep):
     async def _run_async(
         self, fn: Callable, path: str, file_data_path: str
     ) -> list[UncompressStepResponse]:
-        file_data = FileData.from_file(path=file_data_path)
+        file_data = file_data_from_file(path=file_data_path)
         fn_kwargs = {"file_data": file_data}
         if not asyncio.iscoroutinefunction(fn):
             new_file_data = fn(**fn_kwargs)

unstructured_ingest/v2/pipeline/steps/upload.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, TypedDict
 
-from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces.file_data import file_data_from_file
 from unstructured_ingest.v2.interfaces.uploader import UploadContent
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.pipeline.interfaces import BatchPipelineStep
@@ -41,14 +41,14 @@ class UploadStep(BatchPipelineStep):
     @instrument(span_name=STEP_ID)
     def _run_batch(self, contents: list[UploadStepContent]) -> None:
         upload_contents = [
-            UploadContent(path=Path(c["path"]), file_data=FileData.from_file(c["file_data_path"]))
+            UploadContent(path=Path(c["path"]), file_data=file_data_from_file(c["file_data_path"]))
             for c in contents
         ]
         self.process.run_batch(contents=upload_contents)
 
     async def _run_async(self, path: str, file_data_path: str, fn: Optional[Callable] = None):
         fn = fn or self.process.run_async
-        fn_kwargs = {"path": Path(path), "file_data": FileData.from_file(path=file_data_path)}
+        fn_kwargs = {"path": Path(path), "file_data": file_data_from_file(path=file_data_path)}
         if not asyncio.iscoroutinefunction(fn):
             fn(**fn_kwargs)
         elif semaphore := self.context.semaphore:

unstructured_ingest/v2/processes/connectors/__init__.py

@@ -40,6 +40,8 @@ from .milvus import CONNECTOR_TYPE as MILVUS_CONNECTOR_TYPE
 from .milvus import milvus_destination_entry
 from .mongodb import CONNECTOR_TYPE as MONGODB_CONNECTOR_TYPE
 from .mongodb import mongodb_destination_entry, mongodb_source_entry
+from .neo4j import CONNECTOR_TYPE as NEO4J_CONNECTOR_TYPE
+from .neo4j import neo4j_destination_entry
 from .onedrive import CONNECTOR_TYPE as ONEDRIVE_CONNECTOR_TYPE
 from .onedrive import onedrive_destination_entry, onedrive_source_entry
 from .outlook import CONNECTOR_TYPE as OUTLOOK_CONNECTOR_TYPE
@@ -74,6 +76,7 @@ add_destination_entry(destination_type=LOCAL_CONNECTOR_TYPE, entry=local_destina
 add_source_entry(source_type=ONEDRIVE_CONNECTOR_TYPE, entry=onedrive_source_entry)
 add_destination_entry(destination_type=ONEDRIVE_CONNECTOR_TYPE, entry=onedrive_destination_entry)
 
+add_destination_entry(destination_type=NEO4J_CONNECTOR_TYPE, entry=neo4j_destination_entry)
 
 add_source_entry(source_type=SALESFORCE_CONNECTOR_TYPE, entry=salesforce_source_entry)
 

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -1,14 +1,11 @@
-import copy
 import csv
 import hashlib
-import json
-import sys
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
 from typing import TYPE_CHECKING, Any, Generator, Optional
 
-from pydantic import Field, Secret
+from pydantic import BaseModel, Field, Secret
 
 from unstructured_ingest import __name__ as integration_name
 from unstructured_ingest.__version__ import __version__ as integration_version
@@ -17,12 +14,14 @@ from unstructured_ingest.error import (
     SourceConnectionError,
     SourceConnectionNetworkError,
 )
-from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.data_prep import batch_generator, get_data
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.utils.string_and_date_utils import truncate_string_bytes
 from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
+    BatchFileData,
+    BatchItem,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
@@ -54,6 +53,15 @@ CONNECTOR_TYPE = "astradb"
 MAX_CONTENT_PARAM_BYTE_SIZE = 8000
 
 
+class AstraDBAdditionalMetadata(BaseModel):
+    collection_name: str
+    keyspace: Optional[str] = None
+
+
+class AstraDBBatchFileData(BatchFileData):
+    additional_metadata: AstraDBAdditionalMetadata
+
+
 class AstraDBAccessConfig(AccessConfig):
     token: str = Field(description="Astra DB Token with access to the database.")
     api_endpoint: str = Field(description="The API endpoint for the Astra DB.")
@@ -180,9 +188,6 @@ class AstraDBIndexer(Indexer):
 
     def _get_doc_ids(self) -> set[str]:
         """Fetches all document ids in an index"""
-        # Initialize set of ids
-        ids = set()
-
         # Get the collection
         collection = self.get_collection()
 
@@ -195,31 +200,26 @@ class AstraDBIndexer(Indexer):
             astra_db_docs.append(result)
 
         # Create file data for each astra record
-        for astra_record in astra_db_docs:
-            ids.add(astra_record["_id"])
+        ids = sorted([astra_record["_id"] for astra_record in astra_db_docs])
 
-        return ids
+        return set(ids)
 
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+    def run(self, **kwargs: Any) -> Generator[AstraDBBatchFileData, None, None]:
         all_ids = self._get_doc_ids()
         ids = list(all_ids)
         id_batches = batch_generator(ids, self.index_config.batch_size)
 
         for batch in id_batches:
-            # Make sure the hash is always a positive number to create identified
-            identified = str(hash(batch) + sys.maxsize + 1)
-            fd = FileData(
-                identifier=identified,
+            fd = AstraDBBatchFileData(
                 connector_type=CONNECTOR_TYPE,
-                doc_type="batch",
                 metadata=FileDataSourceMetadata(
                     date_processed=str(time()),
                 ),
-                additional_metadata={
-                    "ids": list(batch),
-                    "collection_name": self.index_config.collection_name,
-                    "keyspace": self.index_config.keyspace,
-                },
+                additional_metadata=AstraDBAdditionalMetadata(
+                    collection_name=self.index_config.collection_name,
+                    keyspace=self.index_config.keyspace,
+                ),
+                batch_items=[BatchItem(identifier=b) for b in batch],
             )
             yield fd
 
@@ -248,7 +248,9 @@ class AstraDBDownloader(Downloader):
             writer.writerow(astra_result.keys())
             writer.writerow(astra_result.values())
 
-    def generate_download_response(self, result: dict, file_data: FileData) -> DownloadResponse:
+    def generate_download_response(
+        self, result: dict, file_data: AstraDBBatchFileData
+    ) -> DownloadResponse:
         record_id = result["_id"]
         filename_id = self.get_identifier(record_id=record_id)
         filename = f"{filename_id}.csv"  # csv to preserve column info
@@ -256,7 +258,7 @@ class AstraDBDownloader(Downloader):
         logger.debug(f"Downloading results from record {record_id} as csv to {download_path}")
         download_path.parent.mkdir(parents=True, exist_ok=True)
         try:
-            self.write_astra_result_to_csv(astra_result=result, download_path=download_path)
+            self.write_astra_result_to_csv(astra_result=result, download_path=str(download_path))
         except Exception as e:
             logger.error(
                 f"failed to download from record {record_id} to {download_path}: {e}",
@@ -265,14 +267,12 @@ class AstraDBDownloader(Downloader):
             raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
 
         # modify input file_data for download_response
-        copied_file_data = copy.deepcopy(file_data)
-        copied_file_data.identifier = filename
-        copied_file_data.doc_type = "file"
-        copied_file_data.metadata.date_processed = str(time())
-        copied_file_data.metadata.record_locator = {"document_id": record_id}
-        copied_file_data.additional_metadata.pop("ids", None)
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = filename
+        cast_file_data.metadata.date_processed = str(time())
+        cast_file_data.metadata.record_locator = {"document_id": record_id}
         return super().generate_download_response(
-            file_data=copied_file_data, download_path=download_path
+            file_data=cast_file_data, download_path=download_path
         )
 
     def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
@@ -280,9 +280,10 @@ class AstraDBDownloader(Downloader):
 
     async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
         # Get metadata from file_data
-        ids: list[str] = file_data.additional_metadata["ids"]
-        collection_name: str = file_data.additional_metadata["collection_name"]
-        keyspace: str = file_data.additional_metadata["keyspace"]
+        astra_file_data = AstraDBBatchFileData.cast(file_data=file_data)
+        ids: list[str] = [item.identifier for item in astra_file_data.batch_items]
+        collection_name: str = astra_file_data.additional_metadata.collection_name
+        keyspace: str = astra_file_data.additional_metadata.keyspace
 
         # Retrieve results from async collection
         download_responses = []
@@ -293,7 +294,7 @@ class AstraDBDownloader(Downloader):
         )
         async for result in async_astra_collection.find({"_id": {"$in": ids}}):
             download_responses.append(
-                self.generate_download_response(result=result, file_data=file_data)
+                self.generate_download_response(result=result, file_data=astra_file_data)
             )
         return download_responses
 
@@ -325,29 +326,6 @@ class AstraDBUploadStager(UploadStager):
             "metadata": element_dict,
         }
 
-    def run(
-        self,
-        elements_filepath: Path,
-        file_data: FileData,
-        output_dir: Path,
-        output_filename: str,
-        **kwargs: Any,
-    ) -> Path:
-        with open(elements_filepath) as elements_file:
-            elements_contents = json.load(elements_file)
-        conformed_elements = []
-        for element in elements_contents:
-            conformed_elements.append(self.conform_dict(element_dict=element, file_data=file_data))
-        output_filename_path = Path(output_filename)
-        if output_filename_path.suffix == ".json":
-            output_path = Path(output_dir) / output_filename_path
-        else:
-            output_path = Path(output_dir) / output_filename_path.with_suffix(".json")
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(output_path, "w") as output_file:
-            json.dump(conformed_elements, output_file, indent=2)
-        return output_path
-
 
 @dataclass
 class AstraDBUploader(Uploader):
@@ -386,11 +364,9 @@ class AstraDBUploader(Uploader):
             f"deleted {delete_resp.deleted_count} records from collection {collection.name}"
         )
 
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        with path.open("r") as file:
-            elements_dict = json.load(file)
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
         logger.info(
-            f"writing {len(elements_dict)} objects to destination "
+            f"writing {len(data)} objects to destination "
             f"collection {self.upload_config.collection_name}"
         )
 
@@ -399,9 +375,13 @@ class AstraDBUploader(Uploader):
 
         self.delete_by_record_id(collection=collection, file_data=file_data)
 
-        for chunk in batch_generator(elements_dict, astra_db_batch_size):
+        for chunk in batch_generator(data, astra_db_batch_size):
             collection.insert_many(chunk)
 
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        data = get_data(path=path)
+        self.run_data(data=data, file_data=file_data, **kwargs)
+
 
 astra_db_source_entry = SourceRegistryEntry(
     indexer=AstraDBIndexer,