unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/processes/connectors/duckdb/duckdb.py

@@ -0,0 +1,130 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.error import DestinationConnectionError, ValueError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Uploader,
+    UploaderConfig,
+    UploadStagerConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.processes.connectors.duckdb.base import BaseDuckDBUploadStager
+from unstructured_ingest.utils.data_prep import get_data_df
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from duckdb import DuckDBPyConnection as DuckDBConnection
+    from pandas import DataFrame
+
+CONNECTOR_TYPE = "duckdb"
+
+
+class DuckDBAccessConfig(AccessConfig):
+    pass
+
+
+class DuckDBConnectionConfig(ConnectionConfig):
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    database: Optional[str] = Field(
+        default=None,
+        description="Database name. Path to the DuckDB .db file. If the file does "
+        "not exist, it will be created at the specified path.",
+    )
+    db_schema: Optional[str] = Field(
+        default="main",
+        description="Schema name. Schema in the database where the elements table is located.",
+    )
+    table: Optional[str] = Field(
+        default="elements",
+        description="Table name. Table name into which the elements data is inserted.",
+    )
+    access_config: Secret[DuckDBAccessConfig] = Field(
+        default=DuckDBAccessConfig(), validate_default=True
+    )
+
+    def __post_init__(self):
+        if self.database is None:
+            raise ValueError(
+                "A DuckDB connection requires a path to a *.db or *.duckdb file "
+                "through the `database` argument"
+            )
+
+    @requires_dependencies(["duckdb"], extras="duckdb")
+    @contextmanager
+    def get_client(self) -> Generator["DuckDBConnection", None, None]:
+        import duckdb
+
+        with duckdb.connect(self.database) as client:
+            yield client
+
+    @contextmanager
+    def get_cursor(self) -> Generator["DuckDBConnection", None, None]:
+        with self.get_client() as client, client.cursor() as cursor:
+            yield cursor
+
+
+class DuckDBUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class DuckDBUploadStager(BaseDuckDBUploadStager):
+    upload_stager_config: DuckDBUploadStagerConfig = field(
+        default_factory=lambda: DuckDBUploadStagerConfig()
+    )
+
+
+class DuckDBUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=50, description="[Not-used] Number of records per batch")
+
+
+@dataclass
+class DuckDBUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: DuckDBUploaderConfig
+    connection_config: DuckDBConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            with self.connection_config.get_cursor() as cursor:
+                cursor.execute("SELECT 1;")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def upload_dataframe(self, df: "DataFrame") -> None:
+        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
+
+        with self.connection_config.get_client() as conn:
+            conn.query(
+                f"INSERT INTO {self.connection_config.db_schema}.{self.connection_config.table} BY NAME SELECT * FROM df"  # noqa: E501
+            )
+
+    @requires_dependencies(["pandas"], extras="duckdb")
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df)
+
+    @requires_dependencies(["pandas"], extras="duckdb")
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        df = get_data_df(path)
+        self.upload_dataframe(df=df)
+
+
+duckdb_destination_entry = DestinationRegistryEntry(
+    connection_config=DuckDBConnectionConfig,
+    uploader=DuckDBUploader,
+    uploader_config=DuckDBUploaderConfig,
+    upload_stager=DuckDBUploadStager,
+    upload_stager_config=DuckDBUploadStagerConfig,
+)

unstructured_ingest/processes/connectors/duckdb/motherduck.py

@@ -0,0 +1,130 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.__version__ import __version__ as unstructured_io_ingest_version
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Uploader,
+    UploaderConfig,
+    UploadStagerConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.processes.connectors.duckdb.base import BaseDuckDBUploadStager
+from unstructured_ingest.utils.data_prep import get_data_df
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from duckdb import DuckDBPyConnection as MotherDuckConnection
+    from pandas import DataFrame
+
+CONNECTOR_TYPE = "motherduck"
+
+
+class MotherDuckAccessConfig(AccessConfig):
+    md_token: str = Field(default=None, description="MotherDuck token")
+
+
+class MotherDuckConnectionConfig(ConnectionConfig):
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    database: str = Field(
+        description="Database name. Name of the MotherDuck database.",
+    )
+    db_schema: Optional[str] = Field(
+        default="main",
+        description="Schema name. Schema in the database where the elements table is located.",
+    )
+    table: Optional[str] = Field(
+        default="elements",
+        description="Table name. Table name into which the elements data is inserted.",
+    )
+    access_config: Secret[MotherDuckAccessConfig] = Field(
+        default=MotherDuckAccessConfig(), validate_default=True
+    )
+
+    @requires_dependencies(["duckdb"], extras="duckdb")
+    @contextmanager
+    def get_client(self) -> Generator["MotherDuckConnection", None, None]:
+        import duckdb
+
+        access_config = self.access_config.get_secret_value()
+        with duckdb.connect(
+            f"md:?motherduck_token={access_config.md_token}",
+            config={
+                "custom_user_agent": f"unstructured-io-ingest/{unstructured_io_ingest_version}"
+            },
+        ) as conn:
+            conn.sql(f'USE "{self.database}"')
+            yield conn
+
+    @contextmanager
+    def get_cursor(self) -> Generator["MotherDuckConnection", None, None]:
+        with self.get_client() as client, client.cursor() as cursor:
+            yield cursor
+
+
+class MotherDuckUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+@dataclass
+class MotherDuckUploadStager(BaseDuckDBUploadStager):
+    upload_stager_config: MotherDuckUploadStagerConfig = field(
+        default_factory=lambda: MotherDuckUploadStagerConfig()
+    )
+
+
+class MotherDuckUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=50, description="[Not-used] Number of records per batch")
+
+
+@dataclass
+class MotherDuckUploader(Uploader):
+    connector_type: str = CONNECTOR_TYPE
+    upload_config: MotherDuckUploaderConfig
+    connection_config: MotherDuckConnectionConfig
+
+    def precheck(self) -> None:
+        try:
+            with self.connection_config.get_cursor() as cursor:
+                cursor.execute("SELECT 1;")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def upload_dataframe(self, df: "DataFrame") -> None:
+        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
+        database = self.connection_config.database
+        db_schema = self.connection_config.db_schema
+        table = self.connection_config.table
+
+        with self.connection_config.get_client() as conn:
+            conn.query(f'INSERT INTO "{database}"."{db_schema}"."{table}" BY NAME SELECT * FROM df')
+
+    @requires_dependencies(["pandas"], extras="duckdb")
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df)
+
+    @requires_dependencies(["pandas"], extras="duckdb")
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        df = get_data_df(path)
+        self.upload_dataframe(df=df)
+
+
+motherduck_destination_entry = DestinationRegistryEntry(
+    connection_config=MotherDuckConnectionConfig,
+    uploader=MotherDuckUploader,
+    uploader_config=MotherDuckUploaderConfig,
+    upload_stager=MotherDuckUploadStager,
+    upload_stager_config=MotherDuckUploadStagerConfig,
+)

unstructured_ingest/processes/connectors/elasticsearch/__init__.py

@@ -0,0 +1,19 @@
+from unstructured_ingest.processes.connector_registry import (
+    add_destination_entry,
+    add_source_entry,
+)
+
+from .elasticsearch import CONNECTOR_TYPE as ELASTICSEARCH_CONNECTOR_TYPE
+from .elasticsearch import elasticsearch_destination_entry, elasticsearch_source_entry
+from .opensearch import CONNECTOR_TYPE as OPENSEARCH_CONNECTOR_TYPE
+from .opensearch import opensearch_destination_entry, opensearch_source_entry
+
+add_source_entry(source_type=ELASTICSEARCH_CONNECTOR_TYPE, entry=elasticsearch_source_entry)
+add_destination_entry(
+    destination_type=ELASTICSEARCH_CONNECTOR_TYPE, entry=elasticsearch_destination_entry
+)
+
+add_source_entry(source_type=OPENSEARCH_CONNECTOR_TYPE, entry=opensearch_source_entry)
+add_destination_entry(
+    destination_type=OPENSEARCH_CONNECTOR_TYPE, entry=opensearch_destination_entry
+)