unstructured-ingest 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/fsspec/fsspec.py

@@ -119,7 +119,7 @@ class FsspecIndexer(Indexer):
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise self.wrap_error(e=e)
 
-    def get_file_data(self) -> list[dict[str, Any]]:
+    def get_file_info(self) -> list[dict[str, Any]]:
         if not self.index_config.recursive:
             # fs.ls does not walk directories
             # directories that are listed in cloud storage can cause problems
@@ -156,24 +156,56 @@ class FsspecIndexer(Indexer):
 
         return random.sample(files, n)
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
         raise NotImplementedError()
 
-    def get_path(self, file_data: dict) -> str:
-        return file_data["name"]
+    def get_path(self, file_info: dict) -> str:
+        return file_info["name"]
 
     def sterilize_info(self, file_data: dict) -> dict:
         return sterilize_dict(data=file_data)
 
+    def create_init_file_data(self, remote_filepath: Optional[str] = None) -> FileData:
+        # Create initial file data that requires no network calls and is constructed purely
+        # with information that exists in the config
+        remote_filepath = remote_filepath or self.index_config.remote_url
+        path_without_protocol = remote_filepath.split("://")[1]
+        rel_path = remote_filepath.replace(path_without_protocol, "").lstrip("/")
+        return FileData(
+            identifier=str(uuid5(NAMESPACE_DNS, remote_filepath)),
+            connector_type=self.connector_type,
+            display_name=remote_filepath,
+            source_identifiers=SourceIdentifiers(
+                filename=Path(remote_filepath).name,
+                rel_path=rel_path or None,
+                fullpath=remote_filepath,
+            ),
+            metadata=FileDataSourceMetadata(url=remote_filepath),
+        )
+
+    def hydrate_file_data(self, init_file_data: FileData):
+        # Get file info
+        with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+            files = client.ls(self.index_config.path_without_protocol, detail=True)
+        filtered_files = [
+            file for file in files if file.get("size") > 0 and file.get("type") == "file"
+        ]
+        if not filtered_files:
+            raise ValueError(f"{init_file_data} did not reference any valid file")
+        if len(filtered_files) > 1:
+            raise ValueError(f"{init_file_data} referenced more than one file")
+        file_info = filtered_files[0]
+        init_file_data.additional_metadata = self.get_metadata(file_info=file_info)
+
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        files = self.get_file_data()
-        for file_data in files:
-            file_path = self.get_path(file_data=file_data)
+        files = self.get_file_info()
+        for file_info in files:
+            file_path = self.get_path(file_info=file_info)
             # Note: we remove any remaining leading slashes (Box introduces these)
             # to get a valid relative path
             rel_path = file_path.replace(self.index_config.path_without_protocol, "").lstrip("/")
 
-            additional_metadata = self.sterilize_info(file_data=file_data)
+            additional_metadata = self.sterilize_info(file_data=file_info)
             additional_metadata["original_file_path"] = file_path
             yield FileData(
                 identifier=str(uuid5(NAMESPACE_DNS, file_path)),
@@ -183,7 +215,7 @@ class FsspecIndexer(Indexer):
                     rel_path=rel_path or None,
                     fullpath=file_path,
                 ),
-                metadata=self.get_metadata(file_data=file_data),
+                metadata=self.get_metadata(file_info=file_info),
                 additional_metadata=additional_metadata,
                 display_name=file_path,
             )

unstructured_ingest/v2/processes/connectors/fsspec/gcs.py

@@ -131,22 +131,22 @@ class GcsIndexer(FsspecIndexer):
     index_config: GcsIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"]
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"]
         date_created = None
         date_modified = None
-        if modified_at_str := file_data.get("updated"):
+        if modified_at_str := file_info.get("updated"):
             date_modified = str(parser.parse(modified_at_str).timestamp())
-        if created_at_str := file_data.get("timeCreated"):
+        if created_at_str := file_info.get("timeCreated"):
             date_created = str(parser.parse(created_at_str).timestamp())
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
-        version = file_data.get("etag")
+        version = file_info.get("etag")
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,
-            "file_id": file_data.get("id"),
+            "file_id": file_info.get("id"),
         }
         return FileDataSourceMetadata(
             date_created=date_created,

unstructured_ingest/v2/processes/connectors/fsspec/s3.py

@@ -110,22 +110,22 @@ class S3Indexer(FsspecIndexer):
     def wrap_error(self, e: Exception) -> Exception:
         return self.connection_config.wrap_error(e=e)
 
-    def get_path(self, file_data: dict) -> str:
-        return file_data["Key"]
+    def get_path(self, file_info: dict) -> str:
+        return file_info["Key"]
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["Key"]
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["Key"]
         date_created = None
         date_modified = None
-        modified = file_data.get("LastModified")
+        modified = file_info.get("LastModified")
         if modified:
             date_created = str(modified.timestamp())
             date_modified = str(modified.timestamp())
 
-        file_size = file_data.get("size") if "size" in file_data else None
-        file_size = file_size or file_data.get("Size")
+        file_size = file_info.get("size") if "size" in file_info else None
+        file_size = file_size or file_info.get("Size")
 
-        version = file_data.get("ETag").rstrip('"').lstrip('"') if "ETag" in file_data else None
+        version = file_info.get("ETag").rstrip('"').lstrip('"') if "ETag" in file_info else None
         metadata: dict[str, str] = {}
         with contextlib.suppress(AttributeError):
             with self.connection_config.get_client(protocol=self.index_config.protocol) as client:

unstructured_ingest/v2/processes/connectors/fsspec/sftp.py

@@ -107,12 +107,12 @@ class SftpIndexer(FsspecIndexer):
             file.identifier = new_identifier
             yield file
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"]
-        date_created = str(file_data.get("time").timestamp()) if "time" in file_data else None
-        date_modified = str(file_data.get("mtime").timestamp()) if "mtime" in file_data else None
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"]
+        date_created = str(file_info.get("time").timestamp()) if "time" in file_info else None
+        date_modified = str(file_info.get("mtime").timestamp()) if "mtime" in file_info else None
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
         record_locator = {
             "protocol": self.index_config.protocol,

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

@@ -15,11 +15,14 @@ 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
+from .vastdb import CONNECTOR_TYPE as VASTDB_CONNECTOR_TYPE
+from .vastdb import vastdb_destination_entry, vastdb_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_source_entry(source_type=VASTDB_CONNECTOR_TYPE, entry=vastdb_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)
@@ -31,3 +34,4 @@ add_destination_entry(
     destination_type=DATABRICKS_DELTA_TABLES_CONNECTOR_TYPE,
     entry=databricks_delta_tables_destination_entry,
 )
+add_destination_entry(destination_type=VASTDB_CONNECTOR_TYPE, entry=vastdb_destination_entry)

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

@@ -3,6 +3,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generator, Optional
 
+import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.v2.logger import logger
@@ -139,7 +140,7 @@ class SingleStoreUploader(SQLUploader):
                 if isinstance(value, (list, dict)):
                     value = json.dumps(value)
                 if column_name in _DATE_COLUMNS:
-                    if value is None:
+                    if value is None or pd.isna(value):
                         parsed.append(None)
                     else:
                         parsed.append(parse_date_string(value))

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

@@ -14,7 +14,7 @@ from dateutil import parser
 from pydantic import BaseModel, Field, Secret
 
 from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
-from unstructured_ingest.utils.data_prep import get_data, get_data_df, split_dataframe
+from unstructured_ingest.utils.data_prep import get_data, get_data_df, split_dataframe, write_data
 from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -314,7 +314,7 @@ class SQLUploadStager(UploadStager):
         output_filename = f"{Path(output_filename).stem}{output_filename_suffix}"
         output_path = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
 
-        self.write_output(output_path=output_path, data=df.to_dict(orient="records"))
+        write_data(path=output_path, data=df.to_dict(orient="records"))
         return output_path
 
 
@@ -332,6 +332,7 @@ class SQLUploader(Uploader):
     upload_config: SQLUploaderConfig
     connection_config: SQLConnectionConfig
     values_delimiter: str = "?"
+    _columns: list[str] = field(init=False, default=None)
 
     def precheck(self) -> None:
         try:
@@ -354,7 +355,7 @@ class SQLUploader(Uploader):
             parsed = []
             for column_name, value in zip(columns, row):
                 if column_name in _DATE_COLUMNS:
-                    if value is None:
+                    if value is None or pd.isna(value):  # pandas is nan
                         parsed.append(None)
                     else:
                         parsed.append(parse_date_string(value))
@@ -364,8 +365,9 @@ class SQLUploader(Uploader):
         return output
 
     def _fit_to_schema(self, df: pd.DataFrame) -> pd.DataFrame:
+        table_columns = self.get_table_columns()
         columns = set(df.columns)
-        schema_fields = set(columns)
+        schema_fields = set(table_columns)
         columns_to_drop = columns - schema_fields
         missing_columns = schema_fields - columns
 
@@ -395,8 +397,8 @@ class SQLUploader(Uploader):
                 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)
-        self._fit_to_schema(df=df)
 
         columns = list(df.columns)
         stmt = "INSERT INTO {table_name} ({columns}) VALUES({values})".format(
@@ -424,9 +426,11 @@ class SQLUploader(Uploader):
                 cursor.executemany(stmt, values)
 
     def get_table_columns(self) -> list[str]:
-        with self.get_cursor() as cursor:
-            cursor.execute(f"SELECT * from {self.upload_config.table_name}")
-            return [desc[0] for desc in cursor.description]
+        if self._columns is None:
+            with self.get_cursor() as cursor:
+                cursor.execute(f"SELECT * from {self.upload_config.table_name} LIMIT 1")
+                self._columns = [desc[0] for desc in cursor.description]
+        return self._columns
 
     def can_delete(self) -> bool:
         return self.upload_config.record_id_key in self.get_table_columns()

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

@@ -4,6 +4,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator
 
+import pandas as pd
 from pydantic import Field, Secret, model_validator
 
 from unstructured_ingest.v2.logger import logger
@@ -141,7 +142,7 @@ class SQLiteUploader(SQLUploader):
                 if isinstance(value, (list, dict)):
                     value = json.dumps(value)
                 if column_name in _DATE_COLUMNS:
-                    if value is None:
+                    if value is None or pd.isna(value):
                         parsed.append(None)
                     else:
                         parsed.append(parse_date_string(value))

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

@@ -0,0 +1,270 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+import numpy as np
+import pandas as pd
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.utils.data_prep import split_dataframe
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
+from unstructured_ingest.v2.interfaces 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 (
+    _COLUMNS,
+    SQLAccessConfig,
+    SqlBatchFileData,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+)
+from unstructured_ingest.v2.utils import get_enhanced_element_id
+
+if TYPE_CHECKING:
+    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:
+            with 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'}",
+    )
+    additional_columns: Optional[list[str]] = Field(
+        default_factory=list, description="Additional columns to include in the upload"
+    )
+
+
+class VastdbUploadStager(SQLUploadStager):
+    upload_stager_config: 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)
+
+        # remove extraneous, not supported columns
+        # but also allow for additional columns
+        approved_columns = set(_COLUMNS).union(self.upload_stager_config.additional_columns)
+        element = {k: v for k, v in data.items() if k in approved_columns}
+        element[RECORD_ID_LABEL] = file_data.identifier
+        return element
+
+    def conform_dataframe(self, df: pd.DataFrame) -> pd.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(["pyarrow"], extras="vastdb")
+    def upload_dataframe(self, df: pd.DataFrame, file_data: FileData) -> None:
+        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,
+)