unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/duckdb/base.py

@@ -0,0 +1,100 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+from unstructured_ingest.utils.data_prep import get_data
+from unstructured_ingest.v2.interfaces import FileData, UploadStager
+from unstructured_ingest.v2.utils import get_enhanced_element_id
+
+_COLUMNS = (
+    "id",
+    "element_id",
+    "text",
+    "embeddings",
+    "type",
+    "system",
+    "layout_width",
+    "layout_height",
+    "points",
+    "url",
+    "version",
+    "date_created",
+    "date_modified",
+    "date_processed",
+    "permissions_data",
+    "record_locator",
+    "category_depth",
+    "parent_id",
+    "attached_filename",
+    "filetype",
+    "last_modified",
+    "file_directory",
+    "filename",
+    "languages",
+    "page_number",
+    "links",
+    "page_name",
+    "link_urls",
+    "link_texts",
+    "sent_from",
+    "sent_to",
+    "subject",
+    "section",
+    "header_footer_type",
+    "emphasized_text_contents",
+    "emphasized_text_tags",
+    "text_as_html",
+    "regex_metadata",
+    "detection_class_prob",
+)
+
+# _DATE_COLUMNS = ("date_created", "date_modified", "date_processed", "last_modified")
+
+
+@dataclass
+class BaseDuckDBUploadStager(UploadStager):
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        data = element_dict.copy()
+        metadata: dict[str, Any] = data.pop("metadata", {})
+        data_source = metadata.pop("data_source", {})
+        coordinates = metadata.pop("coordinates", {})
+
+        data.update(metadata)
+        data.update(data_source)
+        data.update(coordinates)
+
+        data["id"] = get_enhanced_element_id(element_dict=data, file_data=file_data)
+
+        # remove extraneous, not supported columns
+        data = {k: v for k, v in data.items() if k in _COLUMNS}
+        return data
+
+    def run(
+        self,
+        elements_filepath: Path,
+        file_data: FileData,
+        output_dir: Path,
+        output_filename: str,
+        **kwargs: Any,
+    ) -> Path:
+        elements_contents = get_data(path=elements_filepath)
+        output_path = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
+
+        output = [
+            self.conform_dict(element_dict=element_dict, file_data=file_data)
+            for element_dict in elements_contents
+        ]
+        df = pd.DataFrame(data=output)
+
+        for column in filter(
+            lambda x: x in df.columns,
+            ("version", "page_number", "regex_metadata"),
+        ):
+            df[column] = df[column].apply(str)
+
+        data = df.to_dict(orient="records")
+        self.write_output(output_path=output_path, data=data)
+        return output_path

unstructured_ingest/v2/processes/connectors/duckdb/duckdb.py

@@ -0,0 +1,127 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+import pandas as pd
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import get_data_df
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.duckdb.base import BaseDuckDBUploadStager
+
+if TYPE_CHECKING:
+    from duckdb import DuckDBPyConnection as DuckDBConnection
+
+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:
+            with 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: pd.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
+            )
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df)
+
+    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/v2/processes/connectors/duckdb/motherduck.py

@@ -0,0 +1,126 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+import pandas as pd
+from pydantic import Field, Secret
+
+from unstructured_ingest.__version__ import __version__ as unstructured_io_ingest_version
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import get_data_df
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connectors.duckdb.base import BaseDuckDBUploadStager
+
+if TYPE_CHECKING:
+    from duckdb import DuckDBPyConnection as MotherDuckConnection
+
+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:
+            with 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: pd.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
+            )
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        df = pd.DataFrame(data=data)
+        self.upload_dataframe(df=df)
+
+    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/v2/processes/connectors/elasticsearch/__init__.py

@@ -0,0 +1,19 @@
+from unstructured_ingest.v2.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
+)