unstructured-ingest 0.0.25__py3-none-any.whl → 0.1.1__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/mongodb.py

@@ -1,26 +1,37 @@
 import json
+import sys
 from dataclasses import dataclass, field
+from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from time import time
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
 from unstructured_ingest.__version__ import __version__ as unstructured_version
-from unstructured_ingest.error import DestinationConnectionError
-from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.utils.data_prep import batch_generator, flatten_dict
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
     FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
     Uploader,
     UploaderConfig,
     UploadStager,
     UploadStagerConfig,
+    download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
+    SourceRegistryEntry,
 )
 
 if TYPE_CHECKING:
@@ -53,6 +64,207 @@ class MongoDBUploadStagerConfig(UploadStagerConfig):
     pass
 
 
+class MongoDBIndexerConfig(IndexerConfig):
+    batch_size: int = Field(default=100, description="Number of records per batch")
+
+
+class MongoDBDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class MongoDBIndexer(Indexer):
+    connection_config: MongoDBConnectionConfig
+    index_config: MongoDBIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self) -> None:
+        """Validates the connection to the MongoDB server."""
+        try:
+            client = self.create_client()
+            client.admin.command("ping")
+        except Exception as e:
+            logger.error(f"Failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"Failed to validate connection: {e}")
+
+    @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
+
+        access_config = self.connection_config.access_config.get_secret_value()
+
+        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),
+            )
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        """Generates FileData objects for each document in the MongoDB collection."""
+        client = self.create_client()
+        database = client[self.connection_config.database]
+        collection = database[self.connection_config.collection]
+
+        # Get list of document IDs
+        ids = collection.distinct("_id")
+        batch_size = self.index_config.batch_size if self.index_config else 100
+
+        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={
+                    "database": self.connection_config.database,
+                    "collection": self.connection_config.collection,
+                },
+            )
+
+            file_data = FileData(
+                identifier=batch_id,
+                doc_type="batch",
+                connector_type=self.connector_type,
+                metadata=metadata,
+                additional_metadata={
+                    "ids": [str(doc_id) for doc_id in id_batch],
+                },
+            )
+            yield file_data
+
+
+@dataclass
+class MongoDBDownloader(Downloader):
+    download_config: MongoDBDownloaderConfig
+    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
+
+        access_config = self.connection_config.access_config.get_secret_value()
+
+        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),
+            )
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["bson"], extras="mongodb")
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        """Fetches the document from MongoDB and writes it to a file."""
+        from bson.errors import InvalidId
+        from bson.objectid import ObjectId
+
+        client = self.create_client()
+        database = client[self.connection_config.database]
+        collection = database[self.connection_config.collection]
+
+        ids = file_data.additional_metadata.get("ids", [])
+        if not ids:
+            raise ValueError("No document IDs provided in additional_metadata")
+
+        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 = FileData(
+                identifier=str(doc_id),
+                connector_type=self.connector_type,
+                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 = FileDataSourceMetadata(
+                date_created=date_created,  # Include date_created here
+                date_processed=str(time()),
+                record_locator={
+                    "database": self.connection_config.database,
+                    "collection": self.connection_config.collection,
+                    "document_id": str(doc_id),
+                },
+            )
+
+            download_response = self.generate_download_response(
+                file_data=individual_file_data, download_path=download_path
+            )
+            download_responses.append(download_response)
+
+        return download_responses
+
+
 @dataclass
 class MongoDBUploadStager(UploadStager):
     upload_stager_config: MongoDBUploadStagerConfig = field(
@@ -138,3 +350,11 @@ mongodb_destination_entry = DestinationRegistryEntry(
     upload_stager=MongoDBUploadStager,
     upload_stager_config=MongoDBUploadStagerConfig,
 )
+
+mongodb_source_entry = SourceRegistryEntry(
+    connection_config=MongoDBConnectionConfig,
+    indexer_config=MongoDBIndexerConfig,
+    indexer=MongoDBIndexer,
+    downloader_config=MongoDBDownloaderConfig,
+    downloader=MongoDBDownloader,
+)

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

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+)
+
+from .postgres import CONNECTOR_TYPE as POSTGRES_CONNECTOR_TYPE
+from .postgres import postgres_destination_entry
+from .sqlite import CONNECTOR_TYPE as SQLITE_CONNECTOR_TYPE
+from .sqlite import sqlite_destination_entry
+
+add_destination_entry(destination_type=SQLITE_CONNECTOR_TYPE, entry=sqlite_destination_entry)
+add_destination_entry(destination_type=POSTGRES_CONNECTOR_TYPE, entry=postgres_destination_entry)

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

@@ -0,0 +1,177 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+import numpy as np
+import pandas as pd
+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
+from unstructured_ingest.v2.processes.connectors.sql.sql import (
+    _DATE_COLUMNS,
+    SQLAccessConfig,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+    parse_date_string,
+)
+
+if TYPE_CHECKING:
+    from psycopg2.extensions import connection as PostgresConnection
+
+CONNECTOR_TYPE = "postgres"
+
+
+class PostgresAccessConfig(SQLAccessConfig):
+    password: Optional[str] = Field(default=None, description="DB password")
+
+
+class PostgresConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[PostgresAccessConfig] = Field(
+        default=PostgresAccessConfig(), validate_default=True
+    )
+    database: Optional[str] = Field(
+        default=None,
+        description="Database name.",
+    )
+    username: Optional[str] = Field(default=None, description="DB username")
+    host: Optional[str] = Field(default=None, description="DB host")
+    port: Optional[int] = Field(default=5432, description="DB host connection port")
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @requires_dependencies(["psycopg2"], extras="postgres")
+    def get_connection(self) -> "PostgresConnection":
+        from psycopg2 import connect
+
+        access_config = self.access_config.get_secret_value()
+        return connect(
+            user=self.username,
+            password=access_config.password,
+            dbname=self.database,
+            host=self.host,
+            port=self.port,
+        )
+
+
+class PostgresIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class PostgresIndexer(SQLIndexer):
+    connection_config: PostgresConnectionConfig
+    index_config: PostgresIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def _get_doc_ids(self) -> list[str]:
+        connection = self.connection_config.get_connection()
+        with connection.cursor() as cursor:
+            cursor.execute(
+                f"SELECT {self.index_config.id_column} FROM {self.index_config.table_name}"
+            )
+            results = cursor.fetchall()
+            ids = [result[0] for result in results]
+            return ids
+
+
+class PostgresDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class PostgresDownloader(SQLDownloader):
+    connection_config: PostgresConnectionConfig
+    download_config: PostgresDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    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"]
+        connection = self.connection_config.get_connection()
+        with connection.cursor() as cursor:
+            fields = ",".join(self.download_config.fields) if self.download_config.fields else "*"
+            query = "SELECT {fields} FROM {table_name} WHERE {id_column} in ({ids})".format(
+                fields=fields,
+                table_name=table_name,
+                id_column=id_column,
+                ids=",".join([str(i) for i in ids]),
+            )
+            logger.debug(f"running query: {query}")
+            cursor.execute(query)
+            rows = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+            return rows, columns
+
+
+class PostgresUploadStagerConfig(SQLUploadStagerConfig):
+    pass
+
+
+class PostgresUploadStager(SQLUploadStager):
+    upload_stager_config: PostgresUploadStagerConfig
+
+
+class PostgresUploaderConfig(SQLUploaderConfig):
+    pass
+
+
+@dataclass
+class PostgresUploader(SQLUploader):
+    upload_config: PostgresUploaderConfig = field(default_factory=PostgresUploaderConfig)
+    connection_config: PostgresConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def prepare_data(
+        self, columns: list[str], data: tuple[tuple[Any, ...], ...]
+    ) -> list[tuple[Any, ...]]:
+        output = []
+        for row in data:
+            parsed = []
+            for column_name, value in zip(columns, row):
+                if column_name in _DATE_COLUMNS:
+                    if value is None:
+                        parsed.append(None)
+                    else:
+                        parsed.append(parse_date_string(value))
+                else:
+                    parsed.append(value)
+            output.append(tuple(parsed))
+        return output
+
+    def upload_contents(self, path: Path) -> None:
+        df = pd.read_json(path, orient="records", lines=True)
+        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
+        df.replace({np.nan: None}, inplace=True)
+
+        columns = tuple(df.columns)
+        stmt = f"INSERT INTO {self.upload_config.table_name} ({','.join(columns)}) \
+                VALUES({','.join(['%s' for x in columns])})"  # noqa E501
+
+        for rows in pd.read_json(
+            path, orient="records", lines=True, chunksize=self.upload_config.batch_size
+        ):
+            with self.connection_config.get_connection() as conn:
+                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
+                with conn.cursor() as cur:
+                    cur.executemany(stmt, values)
+
+                conn.commit()
+
+
+postgres_destination_entry = DestinationRegistryEntry(
+    connection_config=PostgresConnectionConfig,
+    uploader=PostgresUploader,
+    uploader_config=PostgresUploaderConfig,
+    upload_stager=PostgresUploadStager,
+    upload_stager_config=PostgresUploadStagerConfig,
+)