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

unstructured_ingest/v2/processes/connectors/fsspec/gcs.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
-from typing import Any, Generator, Optional, Union
+from typing import TYPE_CHECKING, Any, Generator, Optional, Union
 
 from dateutil import parser
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.utils.string_and_date_utils import json_to_dict
-from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, FileDataSourceMetadata
+from unstructured_ingest.v2.interfaces import FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -26,6 +27,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
     FsspecUploaderConfig,
 )
 
+if TYPE_CHECKING:
+    from gcsfs import GCSFileSystem
+
 CONNECTOR_TYPE = "gcs"
 
 
@@ -93,6 +97,12 @@ class GcsConnectionConfig(FsspecConnectionConfig):
     access_config: Secret[GcsAccessConfig] = Field(default=GcsAccessConfig(), validate_default=True)
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
+    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["GCSFileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class GcsIndexer(FsspecIndexer):
@@ -100,14 +110,6 @@ class GcsIndexer(FsspecIndexer):
     index_config: GcsIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def precheck(self) -> None:
-        super().precheck()
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = None
@@ -147,14 +149,6 @@ class GcsDownloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[GcsDownloaderConfig] = field(default_factory=GcsDownloaderConfig)
 
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class GcsUploaderConfig(FsspecUploaderConfig):
     pass
@@ -166,22 +160,6 @@ class GcsUploader(FsspecUploader):
     connection_config: GcsConnectionConfig
     upload_config: GcsUploaderConfig = field(default=None)
 
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 gcs_source_entry = SourceRegistryEntry(
     indexer=GcsIndexer,

unstructured_ingest/v2/processes/connectors/fsspec/s3.py

@@ -1,15 +1,13 @@
 import contextlib
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from pathlib import Path
 from time import time
-from typing import Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
-    DownloadResponse,
-    FileData,
     FileDataSourceMetadata,
 )
 from unstructured_ingest.v2.processes.connector_registry import (
@@ -29,6 +27,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
 
 CONNECTOR_TYPE = "s3"
 
+if TYPE_CHECKING:
+    from s3fs import S3FileSystem
+
 
 class S3IndexerConfig(FsspecIndexerConfig):
     pass
@@ -72,6 +73,12 @@ class S3ConnectionConfig(FsspecConnectionConfig):
         )
         return access_configs
 
+    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["S3FileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class S3Indexer(FsspecIndexer):
@@ -97,7 +104,8 @@ class S3Indexer(FsspecIndexer):
         version = file_data.get("ETag").rstrip('"').lstrip('"') if "ETag" in file_data else None
         metadata: dict[str, str] = {}
         with contextlib.suppress(AttributeError):
-            metadata = self.fs.metadata(path=path)
+            with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+                metadata = client.metadata(path=path)
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,
@@ -114,14 +122,6 @@ class S3Indexer(FsspecIndexer):
             filesize_bytes=file_size,
         )
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def precheck(self) -> None:
-        super().precheck()
-
 
 class S3DownloaderConfig(FsspecDownloaderConfig):
     pass
@@ -134,14 +134,6 @@ class S3Downloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[S3DownloaderConfig] = field(default_factory=S3DownloaderConfig)
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class S3UploaderConfig(FsspecUploaderConfig):
     pass
@@ -153,22 +145,6 @@ class S3Uploader(FsspecUploader):
     connection_config: S3ConnectionConfig
     upload_config: S3UploaderConfig = field(default=None)
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 s3_source_entry = SourceRegistryEntry(
     indexer=S3Indexer,

unstructured_ingest/v2/processes/connectors/fsspec/sftp.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
 import os
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from time import time
-from typing import Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 from urllib.parse import urlparse
 
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, FileDataSourceMetadata
+from unstructured_ingest.v2.interfaces import FileData, FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -26,6 +27,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
     FsspecUploaderConfig,
 )
 
+if TYPE_CHECKING:
+    from fsspec.implementations.sftp import SFTPFileSystem
+
 CONNECTOR_TYPE = "sftp"
 
 
@@ -67,6 +71,19 @@ class SftpConnectionConfig(FsspecConnectionConfig):
         }
         return access_config
 
+    @contextmanager
+    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
+    def get_client(self, protocol: str) -> Generator["SFTPFileSystem", None, None]:
+        # The paramiko.SSHClient() client that's opened by the SFTPFileSystem
+        # never gets closed so explicitly adding that as part of this context manager
+        from fsspec import get_filesystem_class
+
+        client: SFTPFileSystem = get_filesystem_class(protocol)(
+            **self.get_access_config(),
+        )
+        yield client
+        client.client.close()
+
 
 @dataclass
 class SftpIndexer(FsspecIndexer):
@@ -74,13 +91,11 @@ class SftpIndexer(FsspecIndexer):
     index_config: SftpIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
     def __post_init__(self):
         parsed_url = urlparse(self.index_config.remote_url)
         self.connection_config.host = parsed_url.hostname or self.connection_config.host
         self.connection_config.port = parsed_url.port or self.connection_config.port
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
         for file in super().run(**kwargs):
             new_identifier = (
@@ -92,10 +107,6 @@ class SftpIndexer(FsspecIndexer):
             file.identifier = new_identifier
             yield file
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def precheck(self) -> None:
-        super().precheck()
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = str(file_data.get("time").timestamp()) if "time" in file_data else None
@@ -128,20 +139,11 @@ class SftpDownloader(FsspecDownloader):
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
     download_config: Optional[SftpDownloaderConfig] = field(default_factory=SftpDownloaderConfig)
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
     def __post_init__(self):
         parsed_url = urlparse(self.download_config.remote_url)
         self.connection_config.host = parsed_url.hostname or self.connection_config.host
         self.connection_config.port = parsed_url.port or self.connection_config.port
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class SftpUploaderConfig(FsspecUploaderConfig):
     pass
@@ -153,22 +155,6 @@ class SftpUploader(FsspecUploader):
     connection_config: SftpConnectionConfig
     upload_config: SftpUploaderConfig = field(default=None)
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 sftp_source_entry = SourceRegistryEntry(
     indexer=SftpIndexer,

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,58 @@ 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
+        filename = f"{doc_id}.txt"
+        file_data.source_identifiers = SourceIdentifiers(
+            filename=filename,
+            fullpath=filename,
+        )
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = str(doc_id)
+
+        # 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 +232,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 = (