unstructured-ingest 1.2.32__py3-none-any.whl

unstructured_ingest/processes/connectors/sql/teradata.py

@@ -0,0 +1,254 @@
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.data_types.file_data import FileData
+from unstructured_ingest.logger import logger
+from unstructured_ingest.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.processes.connectors.sql.sql import (
+    SQLAccessConfig,
+    SqlBatchFileData,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+)
+from unstructured_ingest.utils.data_prep import split_dataframe
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+    from teradatasql import TeradataConnection, TeradataCursor
+
+CONNECTOR_TYPE = "teradata"
+
+
+class TeradataAccessConfig(SQLAccessConfig):
+    password: str = Field(description="Teradata user password")
+
+
+class TeradataConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[TeradataAccessConfig]
+    host: str = Field(description="Teradata server hostname or IP address")
+    user: str = Field(description="Teradata database username")
+    database: Optional[str] = Field(
+        default=None,
+        description="Default database/schema to use for queries",
+    )
+    dbs_port: int = Field(
+        default=1025,
+        description="Teradata database port (default: 1025)",
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @contextmanager
+    @requires_dependencies(["teradatasql"], extras="teradata")
+    def get_connection(self) -> Generator["TeradataConnection", None, None]:
+        from teradatasql import connect
+
+        conn_params = {
+            "host": self.host,
+            "user": self.user,
+            "password": self.access_config.get_secret_value().password,
+            "dbs_port": self.dbs_port,
+            "charset": "UTF8",
+        }
+        if self.database:
+            conn_params["database"] = self.database
+
+        connection = connect(**conn_params)
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["TeradataCursor", None, None]:
+        with self.get_connection() as connection:
+            cursor = connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
+
+
+class TeradataIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class TeradataIndexer(SQLIndexer):
+    connection_config: TeradataConnectionConfig
+    index_config: TeradataIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def _get_doc_ids(self) -> list[str]:
+        """Override to quote identifiers for Teradata reserved word handling."""
+        with self.get_cursor() as cursor:
+            cursor.execute(
+                f'SELECT "{self.index_config.id_column}" FROM "{self.index_config.table_name}"'
+            )
+            results = cursor.fetchall()
+            ids = sorted([result[0] for result in results])
+            return ids
+
+
+class TeradataDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class TeradataDownloader(SQLDownloader):
+    connection_config: TeradataConnectionConfig
+    download_config: TeradataDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "?"
+
+    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:
+            if self.download_config.fields:
+                fields = ",".join([f'"{field}"' for field in self.download_config.fields])
+            else:
+                fields = "*"
+
+            placeholders = ",".join([self.values_delimiter for _ in ids])
+            query = f'SELECT {fields} FROM "{table_name}" WHERE "{id_column}" IN ({placeholders})'
+
+            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 TeradataUploadStagerConfig(SQLUploadStagerConfig):
+    pass
+
+
+@dataclass
+class TeradataUploadStager(SQLUploadStager):
+    upload_stager_config: TeradataUploadStagerConfig = field(
+        default_factory=TeradataUploadStagerConfig
+    )
+
+    def conform_dataframe(self, df: "DataFrame") -> "DataFrame":
+        df = super().conform_dataframe(df)
+
+        # teradatasql driver cannot handle Python lists/dicts, convert to JSON strings
+        # Check a sample of values to detect columns with complex types (10 rows)
+        for column in df.columns:
+            sample = df[column].dropna().head(10)
+            
+            if len(sample) > 0:
+                has_complex_type = sample.apply(
+                    lambda x: isinstance(x, (list, dict))
+                ).any()
+                
+                if has_complex_type:
+                    df[column] = df[column].apply(
+                        lambda x: json.dumps(x) if isinstance(x, (list, dict)) else x
+                    )
+
+        return df
+
+
+class TeradataUploaderConfig(SQLUploaderConfig):
+    pass
+
+
+@dataclass
+class TeradataUploader(SQLUploader):
+    upload_config: TeradataUploaderConfig = field(default_factory=TeradataUploaderConfig)
+    connection_config: TeradataConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "?"
+
+    def get_table_columns(self) -> list[str]:
+        if self._columns is None:
+            with self.get_cursor() as cursor:
+                cursor.execute(f'SELECT TOP 1 * FROM "{self.upload_config.table_name}"')
+                self._columns = [desc[0] for desc in cursor.description]
+        return self._columns
+
+    def delete_by_record_id(self, file_data: FileData) -> None:
+        logger.debug(
+            f"deleting any content with data "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from table {self.upload_config.table_name}"
+        )
+        stmt = (
+            f'DELETE FROM "{self.upload_config.table_name}" '
+            f'WHERE "{self.upload_config.record_id_key}" = {self.values_delimiter}'
+        )
+        with self.get_cursor() as cursor:
+            cursor.execute(stmt, [file_data.identifier])
+            rowcount = cursor.rowcount
+            if rowcount > 0:
+                logger.info(f"deleted {rowcount} rows from table {self.upload_config.table_name}")
+
+    def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
+        import numpy as np
+
+        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 = self._fit_to_schema(df=df)
+        df.replace({np.nan: None}, inplace=True)
+
+        columns = list(df.columns)
+        quoted_columns = [f'"{col}"' for col in columns]
+
+        stmt = "INSERT INTO {table_name} ({columns}) VALUES({values})".format(
+            table_name=f'"{self.upload_config.table_name}"',
+            columns=",".join(quoted_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.get_cursor() as cursor:
+                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
+                logger.debug(f"running query: {stmt}")
+                cursor.executemany(stmt, values)
+
+
+teradata_source_entry = SourceRegistryEntry(
+    connection_config=TeradataConnectionConfig,
+    indexer_config=TeradataIndexerConfig,
+    indexer=TeradataIndexer,
+    downloader_config=TeradataDownloaderConfig,
+    downloader=TeradataDownloader,
+)
+
+teradata_destination_entry = DestinationRegistryEntry(
+    connection_config=TeradataConnectionConfig,
+    uploader=TeradataUploader,
+    uploader_config=TeradataUploaderConfig,
+    upload_stager=TeradataUploadStager,
+    upload_stager_config=TeradataUploadStagerConfig,
+)

unstructured_ingest/processes/connectors/sql/vastdb.py

@@ -0,0 +1,263 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.data_types.file_data import (
+    FileData,
+)
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.logger import logger
+from unstructured_ingest.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.processes.connectors.sql.sql import (
+    SQLAccessConfig,
+    SqlBatchFileData,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+)
+from unstructured_ingest.utils.constants import RECORD_ID_LABEL
+from unstructured_ingest.utils.data_prep import get_enhanced_element_id, split_dataframe
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+    from vastdb import connect as VastdbConnect
+    from vastdb import transaction as VastdbTransaction
+    from vastdb.table import Table as VastdbTable
+
+CONNECTOR_TYPE = "vastdb"
+
+
+class VastdbAccessConfig(SQLAccessConfig):
+    endpoint: Optional[str] = Field(default=None, description="DB endpoint")
+    access_key_id: Optional[str] = Field(default=None, description="access key id")
+    access_key_secret: Optional[str] = Field(default=None, description="access key secret")
+
+
+class VastdbConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[VastdbAccessConfig] = Field(
+        default=VastdbAccessConfig(), validate_default=True
+    )
+    vastdb_bucket: str
+    vastdb_schema: str
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @requires_dependencies(["vastdb"], extras="vastdb")
+    @contextmanager
+    def get_connection(self) -> "VastdbConnect":
+        from vastdb import connect
+
+        access_config = self.access_config.get_secret_value()
+        connection = connect(
+            endpoint=access_config.endpoint,
+            access=access_config.access_key_id,
+            secret=access_config.access_key_secret,
+        )
+        yield connection
+
+    @contextmanager
+    def get_cursor(self) -> "VastdbTransaction":
+        with self.get_connection() as connection, connection.transaction() as transaction:
+            yield transaction
+
+    @contextmanager
+    def get_table(self, table_name: str) -> "VastdbTable":
+        with self.get_cursor() as cursor:
+            bucket = cursor.bucket(self.vastdb_bucket)
+            schema = bucket.schema(self.vastdb_schema)
+            table = schema.table(table_name)
+            yield table
+
+
+class VastdbIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class VastdbIndexer(SQLIndexer):
+    connection_config: VastdbConnectionConfig
+    index_config: VastdbIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def _get_doc_ids(self) -> list[str]:
+        with self.connection_config.get_table(self.index_config.table_name) as table:
+            reader = table.select(columns=[self.index_config.id_column])
+            results = reader.read_all()  # Build a PyArrow Table from the RecordBatchReader
+            ids = sorted([result[self.index_config.id_column] for result in results.to_pylist()])
+            return ids
+
+    def precheck(self) -> None:
+        try:
+            with self.connection_config.get_table(self.index_config.table_name) as table:
+                table.select()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+
+class VastdbDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class VastdbDownloader(SQLDownloader):
+    connection_config: VastdbConnectionConfig
+    download_config: VastdbDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    @requires_dependencies(["ibis"], extras="vastdb")
+    def query_db(self, file_data: SqlBatchFileData) -> tuple[list[tuple], list[str]]:
+        from ibis import _  # imports the Ibis deferred expression
+
+        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_table(table_name) as table:
+            predicate = _[id_column].isin(ids)
+
+            if self.download_config.fields:
+                # Vastdb requires the id column to be included in the fields
+                fields = self.download_config.fields + [id_column]
+                # dict.fromkeys to remove duplicates and keep order
+                reader = table.select(columns=list(dict.fromkeys(fields)), predicate=predicate)
+            else:
+                reader = table.select(predicate=predicate)
+            results = reader.read_all()
+            df = results.to_pandas()
+            return [tuple(r) for r in df.to_numpy()], results.column_names
+
+
+class VastdbUploadStagerConfig(SQLUploadStagerConfig):
+    rename_columns_map: Optional[dict] = Field(
+        default=None,
+        description="Map of column names to rename, ex: {'old_name': 'new_name'}",
+    )
+
+
+@dataclass
+class VastdbUploadStager(SQLUploadStager):
+    upload_stager_config: VastdbUploadStagerConfig = field(default_factory=VastdbUploadStagerConfig)
+
+    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)
+        data[RECORD_ID_LABEL] = file_data.identifier
+        return data
+
+    def conform_dataframe(self, df: "DataFrame") -> "DataFrame":
+        df = super().conform_dataframe(df=df)
+        if self.upload_stager_config.rename_columns_map:
+            df.rename(columns=self.upload_stager_config.rename_columns_map, inplace=True)
+        return df
+
+
+class VastdbUploaderConfig(SQLUploaderConfig):
+    pass
+
+
+@dataclass
+class VastdbUploader(SQLUploader):
+    upload_config: VastdbUploaderConfig = field(default_factory=VastdbUploaderConfig)
+    connection_config: VastdbConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self) -> None:
+        try:
+            with self.connection_config.get_table(self.upload_config.table_name) as table:
+                table.select()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    @requires_dependencies(["pandas"], extras="vastdb")
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        super().run(path=path, file_data=file_data, **kwargs)
+
+    @requires_dependencies(["pyarrow", "pandas"], extras="vastdb")
+    def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
+        import numpy as np
+        import pyarrow as pa
+
+        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)
+        df = self._fit_to_schema(df=df)
+
+        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_table(self.upload_config.table_name) as table:
+                pa_table = pa.Table.from_pandas(rows)
+                table.insert(pa_table)
+
+    def get_table_columns(self) -> list[str]:
+        if self._columns is None:
+            with self.connection_config.get_table(self.upload_config.table_name) as table:
+                self._columns = table.columns().names
+        return self._columns
+
+    @requires_dependencies(["ibis"], extras="vastdb")
+    def delete_by_record_id(self, file_data: FileData) -> None:
+        from ibis import _  # imports the Ibis deferred expression
+
+        logger.debug(
+            f"deleting any content with data "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from table {self.upload_config.table_name}"
+        )
+        predicate = _[self.upload_config.record_id_key].isin([file_data.identifier])
+        with self.connection_config.get_table(self.upload_config.table_name) as table:
+            # Get the internal row id
+            rows_to_delete = table.select(
+                columns=[], predicate=predicate, internal_row_id=True
+            ).read_all()
+            table.delete(rows_to_delete)
+
+
+vastdb_source_entry = SourceRegistryEntry(
+    connection_config=VastdbConnectionConfig,
+    indexer_config=VastdbIndexerConfig,
+    indexer=VastdbIndexer,
+    downloader_config=VastdbDownloaderConfig,
+    downloader=VastdbDownloader,
+)
+
+vastdb_destination_entry = DestinationRegistryEntry(
+    connection_config=VastdbConnectionConfig,
+    uploader=VastdbUploader,
+    uploader_config=VastdbUploaderConfig,
+    upload_stager=VastdbUploadStager,
+    upload_stager_config=VastdbUploadStagerConfig,
+)

unstructured_ingest/processes/connectors/utils.py

@@ -0,0 +1,60 @@
+import json
+from datetime import datetime
+from typing import Any, Union
+
+from dateutil import parser
+from pydantic import ValidationError
+
+from unstructured_ingest.utils.chunking import elements_from_base64_gzipped_json
+
+
+def parse_datetime(date_value: Union[int, str, float, datetime]) -> datetime:
+    if isinstance(date_value, datetime):
+        return date_value
+    elif isinstance(date_value, float):
+        return datetime.fromtimestamp(date_value)
+    elif isinstance(date_value, int):
+        return datetime.fromtimestamp(date_value / 1000)
+
+    try:
+        timestamp = float(date_value)
+        return datetime.fromtimestamp(timestamp)
+    except ValueError:
+        return parser.parse(date_value)
+
+
+def conform_string_to_dict(value: Any) -> dict:
+    if isinstance(value, dict):
+        return value
+    if isinstance(value, str):
+        return json.loads(value)
+    raise ValidationError(f"Input could not be mapped to a valid dict: {value}")
+
+
+def format_and_truncate_orig_elements(
+    element: dict, include_text: bool = False
+) -> list[dict[str, Any]]:
+    """
+    This function is used to format and truncate the orig_elements field in the metadata.
+    This is used to remove the text field and other larger fields from the orig_elements
+    that are not helpful in filtering/searching when used along with chunked elements.
+    """
+    metadata = element.get("metadata", {})
+    raw_orig_elements = metadata.get("orig_elements", None)
+    orig_elements = []
+    if raw_orig_elements is not None:
+        for element in elements_from_base64_gzipped_json(raw_orig_elements):
+            if not include_text:
+                element.pop("text", None)
+            for prop in (
+                "image_base64",
+                "text_as_html",
+                "table_as_cells",
+                "link_urls",
+                "link_texts",
+                "link_start_indexes",
+                "emphasized_text_contents",
+            ):
+                element["metadata"].pop(prop, None)
+            orig_elements.append(element)
+    return orig_elements