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

unstructured_ingest/v2/processes/connectors/mongodb.py

@@ -1,11 +1,10 @@
-import sys
 from contextlib import contextmanager
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 from datetime import datetime
 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.__version__ import __version__ as unstructured_version
 from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
@@ -14,9 +13,12 @@ from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
+    BatchFileData,
+    BatchItem,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     FileData,
     FileDataSourceMetadata,
     Indexer,
@@ -40,6 +42,15 @@ CONNECTOR_TYPE = "mongodb"
 SERVER_API_VERSION = "1"
 
 
+class MongoDBAdditionalMetadata(BaseModel):
+    database: str
+    collection: str
+
+
+class MongoDBBatchFileData(BatchFileData):
+    additional_metadata: MongoDBAdditionalMetadata
+
+
 class MongoDBAccessConfig(AccessConfig):
     uri: Optional[str] = Field(default=None, description="URI to user when connecting")
 
@@ -122,7 +133,7 @@ class MongoDBIndexer(Indexer):
             logger.error(f"Failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"Failed to validate connection: {e}")
 
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+    def run(self, **kwargs: Any) -> Generator[BatchFileData, None, None]:
         """Generates FileData objects for each document in the MongoDB collection."""
         with self.connection_config.get_client() as client:
             database = client[self.index_config.database]
@@ -130,12 +141,12 @@ class MongoDBIndexer(Indexer):
 
             # Get list of document IDs
             ids = collection.distinct("_id")
-            batch_size = self.index_config.batch_size if self.index_config else 100
+
+        ids = sorted(ids)
+        batch_size = self.index_config.batch_size
 
         for id_batch in batch_generator(ids, batch_size=batch_size):
             # Make sure the hash is always a positive number to create identifier
-            batch_id = str(hash(frozenset(id_batch)) + sys.maxsize + 1)
-
             metadata = FileDataSourceMetadata(
                 date_processed=str(time()),
                 record_locator={
@@ -144,14 +155,13 @@ class MongoDBIndexer(Indexer):
                 },
             )
 
-            file_data = FileData(
-                identifier=batch_id,
-                doc_type="batch",
+            file_data = MongoDBBatchFileData(
                 connector_type=self.connector_type,
                 metadata=metadata,
-                additional_metadata={
-                    "ids": [str(doc_id) for doc_id in id_batch],
-                },
+                batch_items=[BatchItem(identifier=str(doc_id)) for doc_id in id_batch],
+                additional_metadata=MongoDBAdditionalMetadata(
+                    collection=self.index_config.collection, database=self.index_config.database
+                ),
             )
             yield file_data
 
@@ -162,26 +172,59 @@ class MongoDBDownloader(Downloader):
     connection_config: MongoDBConnectionConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["pymongo"], extras="mongodb")
-    def create_client(self) -> "MongoClient":
-        from pymongo import MongoClient
-        from pymongo.driver_info import DriverInfo
-        from pymongo.server_api import ServerApi
+    def generate_download_response(
+        self, doc: dict, file_data: MongoDBBatchFileData
+    ) -> DownloadResponse:
+        from bson.objectid import ObjectId
 
-        access_config = self.connection_config.access_config.get_secret_value()
+        doc_id = doc["_id"]
+        doc.pop("_id", None)
 
-        if access_config.uri:
-            return MongoClient(
-                access_config.uri,
-                server_api=ServerApi(version=SERVER_API_VERSION),
-                driver=DriverInfo(name="unstructured", version=unstructured_version),
-            )
-        else:
-            return MongoClient(
-                host=self.connection_config.host,
-                port=self.connection_config.port,
-                server_api=ServerApi(version=SERVER_API_VERSION),
-            )
+        # Extract date_created from the document or ObjectId
+        date_created = None
+        if "date_created" in doc:
+            # If the document has a 'date_created' field, use it
+            date_created = doc["date_created"]
+            if isinstance(date_created, datetime):
+                date_created = date_created.isoformat()
+            else:
+                # Convert to ISO format if it's a string
+                date_created = str(date_created)
+        elif isinstance(doc_id, ObjectId):
+            # Use the ObjectId's generation time
+            date_created = doc_id.generation_time.isoformat()
+
+        flattened_dict = flatten_dict(dictionary=doc)
+        concatenated_values = "\n".join(str(value) for value in flattened_dict.values())
+
+        # Create a FileData object for each document with source_identifiers
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = str(doc_id)
+        filename = f"{doc_id}.txt"
+        cast_file_data.source_identifiers = SourceIdentifiers(
+            filename=filename,
+            fullpath=filename,
+            rel_path=filename,
+        )
+
+        # Determine the download path
+        download_path = self.get_download_path(file_data=cast_file_data)
+        if download_path is None:
+            raise ValueError("Download path could not be determined")
+
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write the concatenated values to the file
+        with open(download_path, "w", encoding="utf8") as f:
+            f.write(concatenated_values)
+
+        # Update metadata
+        cast_file_data.metadata.record_locator["document_id"] = str(doc_id)
+        cast_file_data.metadata.date_created = date_created
+
+        return super().generate_download_response(
+            file_data=cast_file_data, download_path=download_path
+        )
 
     @SourceConnectionError.wrap
     @requires_dependencies(["bson"], extras="mongodb")
@@ -190,82 +233,34 @@ class MongoDBDownloader(Downloader):
         from bson.errors import InvalidId
         from bson.objectid import ObjectId
 
-        client = self.create_client()
-        database = client[file_data.metadata.record_locator["database"]]
-        collection = database[file_data.metadata.record_locator["collection"]]
+        mongo_file_data = MongoDBBatchFileData.cast(file_data=file_data)
 
-        ids = file_data.additional_metadata.get("ids", [])
-        if not ids:
-            raise ValueError("No document IDs provided in additional_metadata")
+        with self.connection_config.get_client() as client:
+            database = client[mongo_file_data.additional_metadata.database]
+            collection = database[mongo_file_data.additional_metadata.collection]
 
-        object_ids = []
-        for doc_id in ids:
-            try:
-                object_ids.append(ObjectId(doc_id))
-            except InvalidId as e:
-                error_message = f"Invalid ObjectId for doc_id '{doc_id}': {str(e)}"
-                logger.error(error_message)
-                raise ValueError(error_message) from e
+            ids = [item.identifier for item in mongo_file_data.batch_items]
 
-        try:
-            docs = list(collection.find({"_id": {"$in": object_ids}}))
-        except Exception as e:
-            logger.error(f"Failed to fetch documents: {e}", exc_info=True)
-            raise e
+            object_ids = []
+            for doc_id in ids:
+                try:
+                    object_ids.append(ObjectId(doc_id))
+                except InvalidId as e:
+                    error_message = f"Invalid ObjectId for doc_id '{doc_id}': {str(e)}"
+                    logger.error(error_message)
+                    raise ValueError(error_message) from e
+
+            try:
+                docs = list(collection.find({"_id": {"$in": object_ids}}))
+            except Exception as e:
+                logger.error(f"Failed to fetch documents: {e}", exc_info=True)
+                raise e
 
         download_responses = []
         for doc in docs:
-            doc_id = doc["_id"]
-            doc.pop("_id", None)
-
-            # Extract date_created from the document or ObjectId
-            date_created = None
-            if "date_created" in doc:
-                # If the document has a 'date_created' field, use it
-                date_created = doc["date_created"]
-                if isinstance(date_created, datetime):
-                    date_created = date_created.isoformat()
-                else:
-                    # Convert to ISO format if it's a string
-                    date_created = str(date_created)
-            elif isinstance(doc_id, ObjectId):
-                # Use the ObjectId's generation time
-                date_created = doc_id.generation_time.isoformat()
-
-            flattened_dict = flatten_dict(dictionary=doc)
-            concatenated_values = "\n".join(str(value) for value in flattened_dict.values())
-
-            # Create a FileData object for each document with source_identifiers
-            individual_file_data = replace(file_data)
-            individual_file_data.identifier = str(doc_id)
-            individual_file_data.source_identifiers = SourceIdentifiers(
-                filename=str(doc_id),
-                fullpath=str(doc_id),
-                rel_path=str(doc_id),
-            )
-
-            # Determine the download path
-            download_path = self.get_download_path(individual_file_data)
-            if download_path is None:
-                raise ValueError("Download path could not be determined")
-
-            download_path.parent.mkdir(parents=True, exist_ok=True)
-            download_path = download_path.with_suffix(".txt")
-
-            # Write the concatenated values to the file
-            with open(download_path, "w", encoding="utf8") as f:
-                f.write(concatenated_values)
-
-            individual_file_data.local_download_path = str(download_path)
-
-            # Update metadata
-            individual_file_data.metadata.record_locator["document_id"] = str(doc_id)
-            individual_file_data.metadata.date_created = date_created
-
-            download_response = self.generate_download_response(
-                file_data=individual_file_data, download_path=download_path
+            download_responses.append(
+                self.generate_download_response(doc=doc, file_data=mongo_file_data)
             )
-            download_responses.append(download_response)
 
         return download_responses
 

unstructured_ingest/v2/processes/connectors/neo4j.py

@@ -10,7 +10,6 @@ from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional
 
-import networkx as nx
 from pydantic import BaseModel, ConfigDict, Field, Secret
 
 from unstructured_ingest.error import DestinationConnectionError
@@ -33,6 +32,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 
 if TYPE_CHECKING:
     from neo4j import AsyncDriver, Auth
+    from networkx import Graph, MultiDiGraph
 
 CONNECTOR_TYPE = "neo4j"
 
@@ -109,7 +109,9 @@ class Neo4jUploadStager(UploadStager):
 
         return output_filepath
 
-    def _create_lexical_graph(self, elements: list[dict], document_node: _Node) -> nx.Graph:
+    def _create_lexical_graph(self, elements: list[dict], document_node: _Node) -> "Graph":
+        import networkx as nx
+
         graph = nx.MultiDiGraph()
         graph.add_node(document_node)
 
@@ -180,7 +182,7 @@ class _GraphData(BaseModel):
     edges: list[_Edge]
 
     @classmethod
-    def from_nx(cls, nx_graph: nx.MultiDiGraph) -> _GraphData:
+    def from_nx(cls, nx_graph: "MultiDiGraph") -> _GraphData:
         nodes = list(nx_graph.nodes())
         edges = [
             _Edge(

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -202,7 +202,7 @@ class OnedriveDownloader(Downloader):
         if file_data.source_identifiers is None or not file_data.source_identifiers.fullpath:
             raise ValueError(
                 f"file data doesn't have enough information to get "
-                f"file content: {file_data.to_dict()}"
+                f"file content: {file_data.model_dump()}"
             )
 
         server_relative_path = file_data.source_identifiers.fullpath

unstructured_ingest/v2/processes/connectors/sql/postgres.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Generator, Optional
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
@@ -13,6 +12,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 )
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
     SQLAccessConfig,
+    SqlBatchFileData,
     SQLConnectionConfig,
     SQLDownloader,
     SQLDownloaderConfig,
@@ -99,12 +99,12 @@ class PostgresDownloader(SQLDownloader):
     connector_type: str = CONNECTOR_TYPE
 
     @requires_dependencies(["psycopg2"], extras="postgres")
-    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
+    def query_db(self, file_data: SqlBatchFileData) -> tuple[list[tuple], list[str]]:
         from psycopg2 import sql
 
-        table_name = file_data.additional_metadata["table_name"]
-        id_column = file_data.additional_metadata["id_column"]
-        ids = tuple(file_data.additional_metadata["ids"])
+        table_name = file_data.additional_metadata.table_name
+        id_column = file_data.additional_metadata.id_column
+        ids = tuple([item.identifier for item in file_data.batch_items])
 
         with self.connection_config.get_cursor() as cursor:
             fields = (

unstructured_ingest/v2/processes/connectors/sql/singlestore.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
-from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
@@ -14,6 +13,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
     _DATE_COLUMNS,
     SQLAccessConfig,
+    SqlBatchFileData,
     SQLConnectionConfig,
     SQLDownloader,
     SQLDownloaderConfig,
@@ -93,10 +93,10 @@ class SingleStoreDownloader(SQLDownloader):
     connector_type: str = CONNECTOR_TYPE
     values_delimiter: str = "%s"
 
-    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
-        table_name = file_data.additional_metadata["table_name"]
-        id_column = file_data.additional_metadata["id_column"]
-        ids = tuple(file_data.additional_metadata["ids"])
+    def query_db(self, file_data: SqlBatchFileData) -> tuple[list[tuple], list[str]]:
+        table_name = file_data.additional_metadata.table_name
+        id_column = file_data.additional_metadata.id_column
+        ids = tuple([item.identifier for item in file_data.batch_items])
         with self.connection_config.get_connection() as sqlite_connection:
             cursor = sqlite_connection.cursor()
             fields = ",".join(self.download_config.fields) if self.download_config.fields else "*"

unstructured_ingest/v2/processes/connectors/sql/snowflake.py

@@ -9,7 +9,6 @@ from pydantic import Field, Secret
 
 from unstructured_ingest.utils.data_prep import split_dataframe
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces.file_data import FileData
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
@@ -17,6 +16,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 )
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
     SQLAccessConfig,
+    SqlBatchFileData,
     SQLConnectionConfig,
     SQLDownloader,
     SQLDownloaderConfig,
@@ -118,10 +118,10 @@ class SnowflakeDownloader(SQLDownloader):
 
     # The actual snowflake module package name is: snowflake-connector-python
     @requires_dependencies(["snowflake"], extras="snowflake")
-    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
-        table_name = file_data.additional_metadata["table_name"]
-        id_column = file_data.additional_metadata["id_column"]
-        ids = file_data.additional_metadata["ids"]
+    def query_db(self, file_data: SqlBatchFileData) -> tuple[list[tuple], list[str]]:
+        table_name = file_data.additional_metadata.table_name
+        id_column = file_data.additional_metadata.id_column
+        ids = [item.identifier for item in file_data.batch_items]
 
         with self.connection_config.get_cursor() as cursor:
             query = "SELECT {fields} FROM {table_name} WHERE {id_column} IN ({values})".format(

unstructured_ingest/v2/processes/connectors/sql/sql.py

@@ -1,9 +1,8 @@
 import hashlib
 import json
-import sys
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field
 from datetime import date, datetime
 from pathlib import Path
 from time import time
@@ -12,13 +11,15 @@ from typing import Any, Generator, Union
 import numpy as np
 import pandas as pd
 from dateutil import parser
-from pydantic import Field, Secret
+from pydantic import BaseModel, Field, Secret
 
 from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
 from unstructured_ingest.utils.data_prep import get_data_df, split_dataframe
 from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
+    BatchFileData,
+    BatchItem,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
@@ -81,6 +82,15 @@ _COLUMNS = (
 _DATE_COLUMNS = ("date_created", "date_modified", "date_processed", "last_modified")
 
 
+class SqlAdditionalMetadata(BaseModel):
+    table_name: str
+    id_column: str
+
+
+class SqlBatchFileData(BatchFileData):
+    additional_metadata: SqlAdditionalMetadata
+
+
 def parse_date_string(date_value: Union[str, int]) -> date:
     try:
         timestamp = float(date_value) / 1000 if isinstance(date_value, int) else float(date_value)
@@ -124,7 +134,7 @@ class SQLIndexer(Indexer, ABC):
                 f"SELECT {self.index_config.id_column} FROM {self.index_config.table_name}"
             )
             results = cursor.fetchall()
-            ids = [result[0] for result in results]
+            ids = sorted([result[0] for result in results])
             return ids
 
     def precheck(self) -> None:
@@ -135,7 +145,7 @@ class SQLIndexer(Indexer, ABC):
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
 
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+    def run(self, **kwargs: Any) -> Generator[SqlBatchFileData, None, None]:
         ids = self._get_doc_ids()
         id_batches: list[frozenset[str]] = [
             frozenset(
@@ -151,19 +161,15 @@ class SQLIndexer(Indexer, ABC):
         ]
         for batch in id_batches:
             # Make sure the hash is always a positive number to create identified
-            identified = str(hash(batch) + sys.maxsize + 1)
-            yield FileData(
-                identifier=identified,
+            yield SqlBatchFileData(
                 connector_type=self.connector_type,
                 metadata=FileDataSourceMetadata(
                     date_processed=str(time()),
                 ),
-                doc_type="batch",
-                additional_metadata={
-                    "ids": list(batch),
-                    "table_name": self.index_config.table_name,
-                    "id_column": self.index_config.id_column,
-                },
+                additional_metadata=SqlAdditionalMetadata(
+                    table_name=self.index_config.table_name, id_column=self.index_config.id_column
+                ),
+                batch_items=[BatchItem(identifier=str(b)) for b in batch],
             )
 
 
@@ -176,7 +182,7 @@ class SQLDownloader(Downloader, ABC):
     download_config: SQLDownloaderConfig
 
     @abstractmethod
-    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
+    def query_db(self, file_data: SqlBatchFileData) -> tuple[list[tuple], list[str]]:
         pass
 
     def sql_to_df(self, rows: list[tuple], columns: list[str]) -> list[pd.DataFrame]:
@@ -185,7 +191,7 @@ class SQLDownloader(Downloader, ABC):
         dfs = [pd.DataFrame([row.values], columns=df.columns) for index, row in df.iterrows()]
         return dfs
 
-    def get_data(self, file_data: FileData) -> list[pd.DataFrame]:
+    def get_data(self, file_data: SqlBatchFileData) -> list[pd.DataFrame]:
         rows, columns = self.query_db(file_data=file_data)
         return self.sql_to_df(rows=rows, columns=columns)
 
@@ -199,10 +205,10 @@ class SQLDownloader(Downloader, ABC):
         return f
 
     def generate_download_response(
-        self, result: pd.DataFrame, file_data: FileData
+        self, result: pd.DataFrame, file_data: SqlBatchFileData
     ) -> DownloadResponse:
-        id_column = file_data.additional_metadata["id_column"]
-        table_name = file_data.additional_metadata["table_name"]
+        id_column = file_data.additional_metadata.id_column
+        table_name = file_data.additional_metadata.table_name
         record_id = result.iloc[0][id_column]
         filename_id = self.get_identifier(table_name=table_name, record_id=record_id)
         filename = f"{filename_id}.csv"
@@ -212,20 +218,19 @@ class SQLDownloader(Downloader, ABC):
         )
         download_path.parent.mkdir(parents=True, exist_ok=True)
         result.to_csv(download_path, index=False)
-        copied_file_data = replace(file_data)
-        copied_file_data.identifier = filename_id
-        copied_file_data.doc_type = "file"
-        copied_file_data.additional_metadata.pop("ids", None)
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = filename_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:
-        data_dfs = self.get_data(file_data=file_data)
+        sql_filedata = SqlBatchFileData.cast(file_data=file_data)
+        data_dfs = self.get_data(file_data=sql_filedata)
         download_responses = []
         for df in data_dfs:
             download_responses.append(
-                self.generate_download_response(result=df, file_data=file_data)
+                self.generate_download_response(result=df, file_data=sql_filedata)
             )
         return download_responses
 

unstructured_ingest/v2/processes/connectors/sql/sqlite.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, Generator
 
 from pydantic import Field, Secret, model_validator
 
-from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
@@ -15,6 +14,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
     _DATE_COLUMNS,
     SQLAccessConfig,
+    SqlBatchFileData,
     SQLConnectionConfig,
     SQLDownloader,
     SQLDownloaderConfig,
@@ -97,10 +97,10 @@ class SQLiteDownloader(SQLDownloader):
     connector_type: str = CONNECTOR_TYPE
     values_delimiter: str = "?"
 
-    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
-        table_name = file_data.additional_metadata["table_name"]
-        id_column = file_data.additional_metadata["id_column"]
-        ids = file_data.additional_metadata["ids"]
+    def query_db(self, file_data: SqlBatchFileData) -> tuple[list[tuple], list[str]]:
+        table_name = file_data.additional_metadata.table_name
+        id_column = file_data.additional_metadata.id_column
+        ids = [item.identifier for item in file_data.batch_items]
         with self.connection_config.get_connection() as sqlite_connection:
             cursor = sqlite_connection.cursor()
             fields = ",".join(self.download_config.fields) if self.download_config.fields else "*"