unstructured-ingest 0.3.13__py3-none-any.whl

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

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from unstructured_ingest.v2.processes.connector_registry import (
+    add_destination_entry,
+    add_source_entry,
+)
+
+from .postgres import CONNECTOR_TYPE as POSTGRES_CONNECTOR_TYPE
+from .postgres import postgres_destination_entry, postgres_source_entry
+from .singlestore import CONNECTOR_TYPE as SINGLESTORE_CONNECTOR_TYPE
+from .singlestore import singlestore_destination_entry, singlestore_source_entry
+from .snowflake import CONNECTOR_TYPE as SNOWFLAKE_CONNECTOR_TYPE
+from .snowflake import snowflake_destination_entry, snowflake_source_entry
+from .sqlite import CONNECTOR_TYPE as SQLITE_CONNECTOR_TYPE
+from .sqlite import sqlite_destination_entry, sqlite_source_entry
+
+add_source_entry(source_type=SQLITE_CONNECTOR_TYPE, entry=sqlite_source_entry)
+add_source_entry(source_type=POSTGRES_CONNECTOR_TYPE, entry=postgres_source_entry)
+add_source_entry(source_type=SNOWFLAKE_CONNECTOR_TYPE, entry=snowflake_source_entry)
+add_source_entry(source_type=SINGLESTORE_CONNECTOR_TYPE, entry=singlestore_source_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)
+add_destination_entry(destination_type=SNOWFLAKE_CONNECTOR_TYPE, entry=snowflake_destination_entry)
+add_destination_entry(
+    destination_type=SINGLESTORE_CONNECTOR_TYPE, entry=singlestore_destination_entry
+)

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

@@ -0,0 +1,162 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+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.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.sql.sql import (
+    SQLAccessConfig,
+    SqlBatchFileData,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+)
+
+if TYPE_CHECKING:
+    from psycopg2.extensions import connection as PostgresConnection
+    from psycopg2.extensions import cursor as PostgresCursor
+
+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)
+
+    @contextmanager
+    @requires_dependencies(["psycopg2"], extras="postgres")
+    def get_connection(self) -> Generator["PostgresConnection", None, None]:
+        from psycopg2 import connect
+
+        access_config = self.access_config.get_secret_value()
+        connection = connect(
+            user=self.username,
+            password=access_config.password,
+            dbname=self.database,
+            host=self.host,
+            port=self.port,
+        )
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["PostgresCursor", None, None]:
+        with self.get_connection() as connection:
+            cursor = connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
+
+
+class PostgresIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class PostgresIndexer(SQLIndexer):
+    connection_config: PostgresConnectionConfig
+    index_config: PostgresIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class PostgresDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class PostgresDownloader(SQLDownloader):
+    connection_config: PostgresConnectionConfig
+    download_config: PostgresDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["psycopg2"], extras="postgres")
+    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([item.identifier for item in file_data.batch_items])
+
+        with self.connection_config.get_cursor() as cursor:
+            fields = (
+                sql.SQL(",").join(sql.Identifier(field) for field in self.download_config.fields)
+                if self.download_config.fields
+                else sql.SQL("*")
+            )
+
+            query = sql.SQL("SELECT {fields} FROM {table_name} WHERE {id_column} IN %s").format(
+                fields=fields,
+                table_name=sql.Identifier(table_name),
+                id_column=sql.Identifier(id_column),
+            )
+            logger.debug(f"running query: {cursor.mogrify(query, (ids,))}")
+            cursor.execute(query, (ids,))
+            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
+    values_delimiter: str = "%s"
+
+
+postgres_source_entry = SourceRegistryEntry(
+    connection_config=PostgresConnectionConfig,
+    indexer_config=PostgresIndexerConfig,
+    indexer=PostgresIndexer,
+    downloader_config=PostgresDownloaderConfig,
+    downloader=PostgresDownloader,
+)
+
+postgres_destination_entry = DestinationRegistryEntry(
+    connection_config=PostgresConnectionConfig,
+    uploader=PostgresUploader,
+    uploader_config=PostgresUploaderConfig,
+    upload_stager=PostgresUploadStager,
+    upload_stager_config=PostgresUploadStagerConfig,
+)

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

@@ -0,0 +1,166 @@
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.sql.sql import (
+    _DATE_COLUMNS,
+    SQLAccessConfig,
+    SqlBatchFileData,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+    parse_date_string,
+)
+
+if TYPE_CHECKING:
+    from singlestoredb.connection import Connection as SingleStoreConnection
+    from singlestoredb.connection import Cursor as SingleStoreCursor
+
+CONNECTOR_TYPE = "singlestore"
+
+
+class SingleStoreAccessConfig(SQLAccessConfig):
+    password: Optional[str] = Field(default=None, description="SingleStore password")
+
+
+class SingleStoreConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[SingleStoreAccessConfig]
+    host: Optional[str] = Field(default=None, description="SingleStore host")
+    port: Optional[int] = Field(default=None, description="SingleStore port")
+    user: Optional[str] = Field(default=None, description="SingleStore user")
+    database: Optional[str] = Field(default=None, description="SingleStore database")
+
+    @contextmanager
+    def get_connection(self) -> Generator["SingleStoreConnection", None, None]:
+        import singlestoredb as s2
+
+        connection = s2.connect(
+            host=self.host,
+            port=self.port,
+            database=self.database,
+            user=self.user,
+            password=self.access_config.get_secret_value().password,
+        )
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["SingleStoreCursor", None, None]:
+        with self.get_connection() as connection:
+            with connection.cursor() as cursor:
+                try:
+                    yield cursor
+                finally:
+                    cursor.close()
+
+
+class SingleStoreIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class SingleStoreIndexer(SQLIndexer):
+    connection_config: SingleStoreConnectionConfig
+    index_config: SingleStoreIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class SingleStoreDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class SingleStoreDownloader(SQLDownloader):
+    connection_config: SingleStoreConnectionConfig
+    download_config: SingleStoreDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "%s"
+
+    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 "*"
+            query = (
+                f"SELECT {fields} FROM {table_name} WHERE {id_column} IN {self.values_delimiter}"
+            )
+            logger.debug(f"running query: {query}\nwith values: {(ids,)}")
+            cursor.execute(query, (ids,))
+            rows = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+            return rows, columns
+
+
+class SingleStoreUploadStagerConfig(SQLUploadStagerConfig):
+    pass
+
+
+class SingleStoreUploadStager(SQLUploadStager):
+    upload_stager_config: SingleStoreUploadStagerConfig
+
+
+class SingleStoreUploaderConfig(SQLUploaderConfig):
+    pass
+
+
+@dataclass
+class SingleStoreUploader(SQLUploader):
+    upload_config: SingleStoreUploaderConfig = field(default_factory=SingleStoreUploaderConfig)
+    connection_config: SingleStoreConnectionConfig
+    values_delimiter: str = "%s"
+    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 isinstance(value, (list, dict)):
+                    value = json.dumps(value)
+                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
+
+
+singlestore_source_entry = SourceRegistryEntry(
+    connection_config=SingleStoreConnectionConfig,
+    indexer_config=SingleStoreIndexerConfig,
+    indexer=SingleStoreIndexer,
+    downloader_config=SingleStoreDownloaderConfig,
+    downloader=SingleStoreDownloader,
+)
+
+singlestore_destination_entry = DestinationRegistryEntry(
+    connection_config=SingleStoreConnectionConfig,
+    uploader=SingleStoreUploader,
+    uploader_config=SingleStoreUploaderConfig,
+    upload_stager=SingleStoreUploadStager,
+    upload_stager_config=SingleStoreUploadStagerConfig,
+)

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

@@ -0,0 +1,210 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Generator, Optional
+
+import numpy as np
+import pandas as pd
+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,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.sql.sql import (
+    SQLAccessConfig,
+    SqlBatchFileData,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+)
+
+if TYPE_CHECKING:
+    from snowflake.connector import SnowflakeConnection
+    from snowflake.connector.cursor import SnowflakeCursor
+
+CONNECTOR_TYPE = "snowflake"
+
+
+class SnowflakeAccessConfig(SQLAccessConfig):
+    password: Optional[str] = Field(default=None, description="DB password")
+
+
+class SnowflakeConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[SnowflakeAccessConfig] = Field(
+        default=SnowflakeAccessConfig(), validate_default=True
+    )
+    account: str = Field(
+        default=None,
+        description="Your account identifier. The account identifier "
+        "does not include the snowflakecomputing.com suffix.",
+    )
+    user: Optional[str] = Field(default=None, description="DB username")
+    host: Optional[str] = Field(default=None, description="DB host")
+    port: Optional[int] = Field(default=443, description="DB host connection port")
+    database: str = Field(
+        default=None,
+        description="Database name.",
+    )
+    db_schema: str = Field(default=None, description="Database schema.", alias="schema")
+    role: str = Field(
+        default=None,
+        description="Database role.",
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @contextmanager
+    # The actual snowflake module package name is: snowflake-connector-python
+    @requires_dependencies(["snowflake"], extras="snowflake")
+    def get_connection(self) -> Generator["SnowflakeConnection", None, None]:
+        # https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#label-snowflake-connector-methods-connect
+        from snowflake.connector import connect
+
+        connect_kwargs = self.model_dump()
+        connect_kwargs["schema"] = connect_kwargs.pop("db_schema")
+        connect_kwargs.pop("access_configs", None)
+        connect_kwargs["password"] = self.access_config.get_secret_value().password
+        # https://peps.python.org/pep-0249/#paramstyle
+        connect_kwargs["paramstyle"] = "qmark"
+        # remove anything that is none
+        active_kwargs = {k: v for k, v in connect_kwargs.items() if v is not None}
+        connection = connect(**active_kwargs)
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["SnowflakeCursor", None, None]:
+        with self.get_connection() as connection:
+            cursor = connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
+
+
+class SnowflakeIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class SnowflakeIndexer(SQLIndexer):
+    connection_config: SnowflakeConnectionConfig
+    index_config: SnowflakeIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class SnowflakeDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class SnowflakeDownloader(SQLDownloader):
+    connection_config: SnowflakeConnectionConfig
+    download_config: SnowflakeDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "?"
+
+    # The actual snowflake module package name is: snowflake-connector-python
+    @requires_dependencies(["snowflake"], extras="snowflake")
+    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(
+                table_name=table_name,
+                id_column=id_column,
+                fields=(
+                    ",".join(self.download_config.fields) if self.download_config.fields else "*"
+                ),
+                values=",".join([self.values_delimiter for _ in ids]),
+            )
+            logger.debug(f"running query: {query}\nwith values: {ids}")
+            cursor.execute(query, ids)
+            rows = [
+                tuple(row.values()) if isinstance(row, dict) else row for row in cursor.fetchall()
+            ]
+            columns = [col[0] for col in cursor.description]
+            return rows, columns
+
+
+class SnowflakeUploadStagerConfig(SQLUploadStagerConfig):
+    pass
+
+
+class SnowflakeUploadStager(SQLUploadStager):
+    upload_stager_config: SnowflakeUploadStagerConfig
+
+
+class SnowflakeUploaderConfig(SQLUploaderConfig):
+    pass
+
+
+@dataclass
+class SnowflakeUploader(SQLUploader):
+    upload_config: SnowflakeUploaderConfig = field(default_factory=SnowflakeUploaderConfig)
+    connection_config: SnowflakeConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "?"
+
+    def upload_dataframe(self, df: pd.DataFrame, file_data: FileData) -> None:
+        if self.can_delete():
+            self.delete_by_record_id(file_data=file_data)
+        else:
+            logger.warning(
+                f"table doesn't contain expected "
+                f"record id column "
+                f"{self.upload_config.record_id_key}, skipping delete"
+            )
+        df.replace({np.nan: None}, inplace=True)
+        self._fit_to_schema(df=df, columns=self.get_table_columns())
+
+        columns = list(df.columns)
+        stmt = "INSERT INTO {table_name} ({columns}) VALUES({values})".format(
+            table_name=self.upload_config.table_name,
+            columns=",".join(columns),
+            values=",".join([self.values_delimiter for _ in columns]),
+        )
+        logger.info(
+            f"writing a total of {len(df)} elements via"
+            f" document batches to destination"
+            f" table named {self.upload_config.table_name}"
+            f" with batch size {self.upload_config.batch_size}"
+        )
+        for rows in split_dataframe(df=df, chunk_size=self.upload_config.batch_size):
+            with self.connection_config.get_cursor() as cursor:
+                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
+                # TODO: executemany break on 'Binding data in type (list) is not supported'
+                for val in values:
+                    logger.debug(f"running query: {stmt}\nwith values: {val}")
+                    cursor.execute(stmt, val)
+
+
+snowflake_source_entry = SourceRegistryEntry(
+    connection_config=SnowflakeConnectionConfig,
+    indexer_config=SnowflakeIndexerConfig,
+    indexer=SnowflakeIndexer,
+    downloader_config=SnowflakeDownloaderConfig,
+    downloader=SnowflakeDownloader,
+)
+
+snowflake_destination_entry = DestinationRegistryEntry(
+    connection_config=SnowflakeConnectionConfig,
+    uploader=SnowflakeUploader,
+    uploader_config=SnowflakeUploaderConfig,
+    upload_stager=SnowflakeUploadStager,
+    upload_stager_config=SnowflakeUploadStagerConfig,
+)