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

unstructured_ingest/v2/processes/connectors/sql.py

@@ -1,275 +0,0 @@
-import json
-import uuid
-from dataclasses import dataclass, field
-from datetime import date, datetime
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
-
-import numpy as np
-import pandas as pd
-from dateutil import parser
-from pydantic import Field, Secret
-
-from unstructured_ingest.error import DestinationConnectionError
-from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import (
-    AccessConfig,
-    ConnectionConfig,
-    FileData,
-    Uploader,
-    UploaderConfig,
-    UploadStager,
-    UploadStagerConfig,
-)
-from unstructured_ingest.v2.logger import logger
-from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
-
-if TYPE_CHECKING:
-    from sqlite3 import Connection as SqliteConnection
-
-    from psycopg2.extensions import connection as PostgresConnection
-
-CONNECTOR_TYPE = "sql"
-ELEMENTS_TABLE_NAME = "elements"
-SQLITE_DB = "sqlite"
-POSTGRESQL_DB = "postgresql"
-
-
-class SQLAccessConfig(AccessConfig):
-    username: Optional[str] = Field(default=None, description="DB username")
-    password: Optional[str] = Field(default=None, description="DB password")
-
-
-class SQLConnectionConfig(ConnectionConfig):
-    db_type: Literal["sqlite", "postgresql"] = Field(
-        default=SQLITE_DB, description="Type of the database backend"
-    )
-    database: Optional[str] = Field(
-        default=None,
-        description="Database name. For sqlite databases, this is the path to the .db file.",
-    )
-    host: Optional[str] = Field(default=None, description="DB host")
-    port: Optional[int] = Field(default=5432, description="DB host connection port")
-    access_config: Secret[SQLAccessConfig] = Field(default=SQLAccessConfig(), validate_default=True)
-    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
-
-    def __post_init__(self):
-        if (self.db_type == SQLITE_DB) and (self.database is None):
-            raise ValueError(
-                "A sqlite connection requires a path to a *.db file "
-                "through the `database` argument"
-            )
-
-
-class SQLUploadStagerConfig(UploadStagerConfig):
-    pass
-
-
-_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")
-
-
-def parse_date_string(date_value: Union[str, int]) -> date:
-    try:
-        timestamp = float(date_value) / 1000 if isinstance(date_value, int) else float(date_value)
-        return datetime.fromtimestamp(timestamp)
-    except Exception as e:
-        logger.debug(f"date {date_value} string not a timestamp: {e}")
-    return parser.parse(date_value)
-
-
-@dataclass
-class SQLUploadStager(UploadStager):
-    upload_stager_config: SQLUploadStagerConfig = field(
-        default_factory=lambda: SQLUploadStagerConfig()
-    )
-
-    def run(
-        self,
-        elements_filepath: Path,
-        file_data: FileData,
-        output_dir: Path,
-        output_filename: str,
-        **kwargs: Any,
-    ) -> Path:
-        with open(elements_filepath) as elements_file:
-            elements_contents: list[dict] = json.load(elements_file)
-        output_path = Path(output_dir) / Path(f"{output_filename}.json")
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-
-        output = []
-        for data in elements_contents:
-            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"] = str(uuid.uuid4())
-
-            # remove extraneous, not supported columns
-            data = {k: v for k, v in data.items() if k in _COLUMNS}
-
-            output.append(data)
-
-        df = pd.DataFrame.from_dict(output)
-        for column in filter(lambda x: x in df.columns, _DATE_COLUMNS):
-            df[column] = df[column].apply(parse_date_string)
-        for column in filter(
-            lambda x: x in df.columns,
-            ("permissions_data", "record_locator", "points", "links"),
-        ):
-            df[column] = df[column].apply(
-                lambda x: json.dumps(x) if isinstance(x, (list, dict)) else None
-            )
-        for column in filter(
-            lambda x: x in df.columns,
-            ("version", "page_number", "regex_metadata"),
-        ):
-            df[column] = df[column].apply(str)
-
-        with output_path.open("w") as output_file:
-            df.to_json(output_file, orient="records", lines=True)
-        return output_path
-
-
-class SQLUploaderConfig(UploaderConfig):
-    batch_size: int = Field(default=50, description="Number of records per batch")
-
-
-@dataclass
-class SQLUploader(Uploader):
-    connector_type: str = CONNECTOR_TYPE
-    upload_config: SQLUploaderConfig
-    connection_config: SQLConnectionConfig
-
-    def precheck(self) -> None:
-        try:
-            cursor = self.connection().cursor()
-            cursor.execute("SELECT 1;")
-            cursor.close()
-        except Exception as e:
-            logger.error(f"failed to validate connection: {e}", exc_info=True)
-            raise DestinationConnectionError(f"failed to validate connection: {e}")
-
-    @property
-    def connection(self) -> Callable[[], Union["SqliteConnection", "PostgresConnection"]]:
-        if self.connection_config.db_type == POSTGRESQL_DB:
-            return self._make_psycopg_connection
-        elif self.connection_config.db_type == SQLITE_DB:
-            return self._make_sqlite_connection
-        raise ValueError(f"Unsupported database {self.connection_config.db_type} connection.")
-
-    def _make_sqlite_connection(self) -> "SqliteConnection":
-        from sqlite3 import connect
-
-        return connect(database=self.connection_config.database)
-
-    @requires_dependencies(["psycopg2"], extras="postgres")
-    def _make_psycopg_connection(self) -> "PostgresConnection":
-        from psycopg2 import connect
-
-        access_config = self.connection_config.access_config.get_secret_value()
-        return connect(
-            user=access_config.username,
-            password=access_config.password,
-            dbname=self.connection_config.database,
-            host=self.connection_config.host,
-            port=self.connection_config.port,
-        )
-
-    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 self.connection_config.db_type == SQLITE_DB and 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
-
-    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 {ELEMENTS_TABLE_NAME} ({','.join(columns)}) \
-                VALUES({','.join(['?' if self.connection_config.db_type==SQLITE_DB else '%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() as conn:
-                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
-                if self.connection_config.db_type == SQLITE_DB:
-                    conn.executemany(stmt, values)
-                else:
-                    with conn.cursor() as cur:
-                        cur.executemany(stmt, values)
-
-                conn.commit()
-
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        self.upload_contents(path=path)
-
-
-sql_destination_entry = DestinationRegistryEntry(
-    connection_config=SQLConnectionConfig,
-    uploader=SQLUploader,
-    uploader_config=SQLUploaderConfig,
-    upload_stager=SQLUploadStager,
-    upload_stager_config=SQLUploadStagerConfig,
-)