unstructured-ingest 0.1.0__tar.gz → 0.1.1__tar.gz

{unstructured_ingest-0.1.0 → unstructured_ingest-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unstructured-ingest
-Version: 0.1.0
+Version: 0.1.1
 Summary: A library that prepares raw documents for downstream ML tasks.
 Home-page: https://github.com/Unstructured-IO/unstructured-ingest
 Author: Unstructured Technologies
@@ -22,13 +22,13 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9.0,<3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
-Requires-Dist: pandas
-Requires-Dist: pydantic>=2.7
 Requires-Dist: tqdm
-Requires-Dist: dataclasses_json
-Requires-Dist: opentelemetry-sdk
 Requires-Dist: python-dateutil
+Requires-Dist: dataclasses_json
+Requires-Dist: pandas
+Requires-Dist: pydantic>=2.7
 Requires-Dist: click
+Requires-Dist: opentelemetry-sdk
 Provides-Extra: remote
 Requires-Dist: unstructured-client>=0.25.8; extra == "remote"
 Provides-Extra: csv
@@ -74,8 +74,8 @@ Provides-Extra: biomed
 Requires-Dist: requests; extra == "biomed"
 Requires-Dist: bs4; extra == "biomed"
 Provides-Extra: box
-Requires-Dist: fsspec; extra == "box"
 Requires-Dist: boxfs; extra == "box"
+Requires-Dist: fsspec; extra == "box"
 Provides-Extra: chroma
 Requires-Dist: chromadb; extra == "chroma"
 Provides-Extra: clarifai
@@ -86,8 +86,8 @@ Requires-Dist: atlassian-python-api; extra == "confluence"
 Provides-Extra: couchbase
 Requires-Dist: couchbase; extra == "couchbase"
 Provides-Extra: delta-table
-Requires-Dist: fsspec; extra == "delta-table"
 Requires-Dist: deltalake; extra == "delta-table"
+Requires-Dist: fsspec; extra == "delta-table"
 Provides-Extra: discord
 Requires-Dist: discord-py; extra == "discord"
 Provides-Extra: dropbox
@@ -96,9 +96,9 @@ Requires-Dist: fsspec; extra == "dropbox"
 Provides-Extra: elasticsearch
 Requires-Dist: elasticsearch[async]; extra == "elasticsearch"
 Provides-Extra: gcs
+Requires-Dist: bs4; extra == "gcs"
 Requires-Dist: gcsfs; extra == "gcs"
 Requires-Dist: fsspec; extra == "gcs"
-Requires-Dist: bs4; extra == "gcs"
 Provides-Extra: github
 Requires-Dist: requests; extra == "github"
 Requires-Dist: pygithub>1.58.0; extra == "github"
@@ -114,25 +114,25 @@ Requires-Dist: atlassian-python-api; extra == "jira"
 Provides-Extra: kafka
 Requires-Dist: confluent-kafka; extra == "kafka"
 Provides-Extra: kdbai
-Requires-Dist: kdbai-client; extra == "kdbai"
+Requires-Dist: kdbai-client>=1.4.0; extra == "kdbai"
 Provides-Extra: milvus
 Requires-Dist: pymilvus; extra == "milvus"
 Provides-Extra: mongodb
 Requires-Dist: pymongo; extra == "mongodb"
 Provides-Extra: notion
 Requires-Dist: notion-client; extra == "notion"
-Requires-Dist: httpx; extra == "notion"
 Requires-Dist: backoff; extra == "notion"
 Requires-Dist: htmlBuilder; extra == "notion"
+Requires-Dist: httpx; extra == "notion"
 Provides-Extra: onedrive
 Requires-Dist: bs4; extra == "onedrive"
-Requires-Dist: msal; extra == "onedrive"
 Requires-Dist: Office365-REST-Python-Client; extra == "onedrive"
+Requires-Dist: msal; extra == "onedrive"
 Provides-Extra: opensearch
 Requires-Dist: opensearch-py; extra == "opensearch"
 Provides-Extra: outlook
-Requires-Dist: msal; extra == "outlook"
 Requires-Dist: Office365-REST-Python-Client; extra == "outlook"
+Requires-Dist: msal; extra == "outlook"
 Provides-Extra: pinecone
 Requires-Dist: pinecone-client>=3.7.1; extra == "pinecone"
 Provides-Extra: postgres
@@ -145,13 +145,13 @@ Provides-Extra: s3
 Requires-Dist: s3fs; extra == "s3"
 Requires-Dist: fsspec; extra == "s3"
 Provides-Extra: sharepoint
-Requires-Dist: msal; extra == "sharepoint"
 Requires-Dist: Office365-REST-Python-Client; extra == "sharepoint"
+Requires-Dist: msal; extra == "sharepoint"
 Provides-Extra: salesforce
 Requires-Dist: simple-salesforce; extra == "salesforce"
 Provides-Extra: sftp
-Requires-Dist: fsspec; extra == "sftp"
 Requires-Dist: paramiko; extra == "sftp"
+Requires-Dist: fsspec; extra == "sftp"
 Provides-Extra: slack
 Requires-Dist: slack_sdk; extra == "slack"
 Provides-Extra: wikipedia

{unstructured_ingest-0.1.0/test/integration/connectors → unstructured_ingest-0.1.1/test/integration/connectors/sql}/test_postgres.py

@@ -1,21 +1,97 @@
 import tempfile
+from contextlib import contextmanager
 from pathlib import Path
 
+import faker
 import pandas as pd
 import pytest
 from psycopg2 import connect
 
-from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG, env_setup_path
 from test.integration.connectors.utils.docker_compose import docker_compose_context
+from test.integration.connectors.utils.validation import (
+    ValidationConfigs,
+    source_connector_validation,
+)
 from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.processes.connectors.sql.postgres import (
     CONNECTOR_TYPE,
     PostgresAccessConfig,
     PostgresConnectionConfig,
+    PostgresDownloader,
+    PostgresDownloaderConfig,
+    PostgresIndexer,
+    PostgresIndexerConfig,
     PostgresUploader,
     PostgresUploadStager,
 )
 
+faker = faker.Faker()
+
+SEED_DATA_ROWS = 40
+
+
+@contextmanager
+def postgres_download_setup() -> None:
+    with docker_compose_context(docker_compose_path=env_setup_path / "sql" / "postgres" / "source"):
+        connection = connect(
+            user="unstructured",
+            password="test",
+            dbname="test_db",
+            host="localhost",
+            port=5433,
+        )
+        with connection.cursor() as cursor:
+            for _ in range(SEED_DATA_ROWS):
+                sql_statment = (
+                    f"INSERT INTO cars (brand, price) VALUES "
+                    f"('{faker.word()}', {faker.random_int()})"
+                )
+                cursor.execute(sql_statment)
+            connection.commit()
+        yield
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+async def test_postgres_source():
+    connect_params = {
+        "host": "localhost",
+        "port": 5433,
+        "database": "test_db",
+        "user": "unstructured",
+        "password": "test",
+    }
+    with postgres_download_setup():
+        with tempfile.TemporaryDirectory() as tmpdir:
+            connection_config = PostgresConnectionConfig(
+                host=connect_params["host"],
+                port=connect_params["port"],
+                database=connect_params["database"],
+                username=connect_params["user"],
+                access_config=PostgresAccessConfig(password=connect_params["password"]),
+            )
+            indexer = PostgresIndexer(
+                connection_config=connection_config,
+                index_config=PostgresIndexerConfig(
+                    table_name="cars", id_column="car_id", batch_size=5
+                ),
+            )
+            downloader = PostgresDownloader(
+                connection_config=connection_config,
+                download_config=PostgresDownloaderConfig(
+                    fields=["car_id", "brand"], download_dir=Path(tmpdir)
+                ),
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="postgres",
+                    expected_num_files=40,
+                ),
+            )
+
 
 def validate_destination(
     connect_params: dict,
@@ -50,7 +126,9 @@ async def test_postgres_destination(upload_file: Path):
     # the postgres destination connector doesn't leverage the file data but is required as an input,
     # mocking it with arbitrary values to meet the base requirements:
     mock_file_data = FileData(identifier="mock file data", connector_type=CONNECTOR_TYPE)
-    with docker_compose_context(docker_compose_path=env_setup_path / "sql"):
+    with docker_compose_context(
+        docker_compose_path=env_setup_path / "sql" / "postgres" / "destination"
+    ):
         with tempfile.TemporaryDirectory() as tmpdir:
             stager = PostgresUploadStager()
             stager_params = {

{unstructured_ingest-0.1.0/test/integration/connectors → unstructured_ingest-0.1.1/test/integration/connectors/sql}/test_sqlite.py

@@ -3,39 +3,99 @@ import tempfile
 from contextlib import contextmanager
 from pathlib import Path
 
+import faker
 import pandas as pd
 import pytest
 
-from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG, env_setup_path
+from test.integration.connectors.utils.validation import (
+    ValidationConfigs,
+    source_connector_validation,
+)
 from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.processes.connectors.sql.sqlite import (
     CONNECTOR_TYPE,
     SQLiteConnectionConfig,
+    SQLiteDownloader,
+    SQLiteDownloaderConfig,
+    SQLiteIndexer,
+    SQLiteIndexerConfig,
     SQLiteUploader,
     SQLiteUploadStager,
 )
 
+faker = faker.Faker()
+
+SEED_DATA_ROWS = 40
+
+
+@contextmanager
+def sqlite_download_setup() -> Path:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = Path(tmpdir) / "mock_database.db"
+        db_init_path = env_setup_path / "sql" / "sqlite" / "source" / "sqlite-schema.sql"
+        assert db_init_path.exists()
+        assert db_init_path.is_file()
+        with sqlite3.connect(database=db_path) as sqlite_connection:
+            cursor = sqlite_connection.cursor()
+            with db_init_path.open("r") as f:
+                query = f.read()
+            cursor.executescript(query)
+            for _ in range(SEED_DATA_ROWS):
+                sql_statment = (
+                    f"INSERT INTO cars (brand, price) "
+                    f"VALUES ('{faker.word()}', {faker.random_int()})"
+                )
+                cursor.execute(sql_statment)
+
+            sqlite_connection.commit()
+            cursor.close()
+        yield db_path
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+async def test_sqlite_source():
+    with sqlite_download_setup() as db_path:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            connection_config = SQLiteConnectionConfig(database_path=db_path)
+            indexer = SQLiteIndexer(
+                connection_config=connection_config,
+                index_config=SQLiteIndexerConfig(
+                    table_name="cars", id_column="car_id", batch_size=5
+                ),
+            )
+            downloader = SQLiteDownloader(
+                connection_config=connection_config,
+                download_config=SQLiteDownloaderConfig(
+                    fields=["car_id", "brand"], download_dir=Path(tmpdir)
+                ),
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="sqlite",
+                    expected_num_files=40,
+                ),
+            )
+
 
 @contextmanager
-def sqlite_setup() -> Path:
+def sqlite_upload_setup() -> Path:
     # Provision the local file that sqlite points to to have the desired schema for the integration
     # tests and make sure the file and connection get cleaned up by using a context manager.
     with tempfile.TemporaryDirectory() as tmpdir:
         db_path = Path(tmpdir) / "elements.db"
-        db_init_path = env_setup_path / "sql" / "sqlite-schema.sql"
+        db_init_path = env_setup_path / "sql" / "sqlite" / "destination" / "sqlite-schema.sql"
         assert db_init_path.exists()
         assert db_init_path.is_file()
-        connection = None
-        try:
-            connection = sqlite3.connect(database=db_path)
+        with sqlite3.connect(database=db_path) as sqlite_connection:
             with db_init_path.open("r") as f:
                 query = f.read()
-            cursor = connection.cursor()
+            cursor = sqlite_connection.cursor()
             cursor.executescript(query)
-            yield db_path
-        finally:
-            if connection:
-                connection.close()
+        yield db_path
 
 
 def validate_destination(db_path: Path, expected_num_elements: int):
@@ -62,7 +122,7 @@ async def test_sqlite_destination(upload_file: Path):
     # the sqlite destination connector doesn't leverage the file data but is required as an input,
     # mocking it with arbitrary values to meet the base requirements:
     mock_file_data = FileData(identifier="mock file data", connector_type=CONNECTOR_TYPE)
-    with sqlite_setup() as db_path:
+    with sqlite_upload_setup() as db_path:
         with tempfile.TemporaryDirectory() as tmpdir:
             stager = SQLiteUploadStager()
             stager_params = {

{unstructured_ingest-0.1.0 → unstructured_ingest-0.1.1}/test/integration/connectors/utils/constants.py

@@ -3,5 +3,5 @@ from pathlib import Path
 SOURCE_TAG = "source"
 DESTINATION_TAG = "destination"
 
-env_setup_path = Path(__file__).parents[4] / "test_e2e" / "env_setup"
+env_setup_path = Path(__file__).parents[1] / "env_setup"
 expected_results_path = Path(__file__).parents[1] / "expected_results"

{unstructured_ingest-0.1.0 → unstructured_ingest-0.1.1}/test/integration/connectors/utils/validation.py

@@ -180,8 +180,13 @@ async def source_connector_validation(
             resp = await downloader.run_async(file_data=file_data)
         else:
             resp = downloader.run(file_data=file_data)
-        postdownload_file_data = replace(resp["file_data"])
-        all_postdownload_file_data.append(postdownload_file_data)
+        if isinstance(resp, list):
+            for r in resp:
+                postdownload_file_data = replace(r["file_data"])
+                all_postdownload_file_data.append(postdownload_file_data)
+        else:
+            postdownload_file_data = replace(resp["file_data"])
+            all_postdownload_file_data.append(postdownload_file_data)
     if not overwrite_fixtures:
         run_all_validations(
             configs=configs,

unstructured_ingest-0.1.1/unstructured_ingest/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.1"  # pragma: no cover

{unstructured_ingest-0.1.0 → unstructured_ingest-0.1.1}/unstructured_ingest/v2/processes/connectors/fsspec/fsspec.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import random
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator, Optional, TypeVar
@@ -63,6 +64,7 @@ class FileConfig(BaseModel):
 
 class FsspecIndexerConfig(FileConfig, IndexerConfig):
     recursive: bool = False
+    sample_n_files: Optional[int] = None
 
 
 class FsspecAccessConfig(AccessConfig):
@@ -128,8 +130,23 @@ class FsspecIndexer(Indexer):
         filtered_files = [
             file for file in files if file.get("size") > 0 and file.get("type") == "file"
         ]
+
+        if self.index_config.sample_n_files:
+            filtered_files = self.sample_n_files(filtered_files, self.index_config.sample_n_files)
+
         return filtered_files
 
+    def sample_n_files(self, files: list[dict[str, Any]], n) -> list[dict[str, Any]]:
+        if len(files) <= n:
+            logger.warning(
+                f"number of files to be sampled={n} is not smaller than the number"
+                f" of files found ({len(files)}). Returning all of the files as the"
+                " sample."
+            )
+            return files
+
+        return random.sample(files, n)
+
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         raise NotImplementedError()
 

{unstructured_ingest-0.1.0 → unstructured_ingest-0.1.1}/unstructured_ingest/v2/processes/connectors/kdbai.py

@@ -26,7 +26,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
 )
 
 if TYPE_CHECKING:
-    from kdbai_client import Session, Table
+    from kdbai_client import Database, Session, Table
 
 CONNECTOR_TYPE = "kdbai"
 
@@ -99,6 +99,9 @@ class KdbaiUploadStager(UploadStager):
 
 
 class KdbaiUploaderConfig(UploaderConfig):
+    database_name: str = Field(
+        default="default", description="The name of the KDBAI database to write into."
+    )
     table_name: str = Field(description="The name of the KDBAI table to write into.")
     batch_size: int = Field(default=100, description="Number of records per batch")
 
@@ -111,24 +114,29 @@ class KdbaiUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            self.get_table()
+            self.get_database()
         except Exception as e:
             logger.error(f"Failed to validate connection {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    def get_table(self) -> "Table":
+    def get_database(self) -> "Database":
         session: Session = self.connection_config.get_session()
-        table = session.table(self.upload_config.table_name)
+        db = session.database(self.upload_config.database_name)
+        return db
+
+    def get_table(self) -> "Table":
+        db = self.get_database()
+        table = db.table(self.upload_config.table_name)
         return table
 
     def upsert_batch(self, batch: pd.DataFrame):
         table = self.get_table()
-        table.insert(data=batch)
+        table.insert(batch)
 
     def process_dataframe(self, df: pd.DataFrame):
         logger.debug(
             f"uploading {len(df)} entries to {self.connection_config.endpoint} "
-            f"db in table {self.upload_config.table_name}"
+            f"db {self.upload_config.database_name} in table {self.upload_config.table_name}"
         )
         for _, batch_df in df.groupby(np.arange(len(df)) // self.upload_config.batch_size):
             self.upsert_batch(batch=batch_df)

{unstructured_ingest-0.1.0 → unstructured_ingest-0.1.1}/unstructured_ingest/v2/processes/connectors/sql/postgres.py

@@ -7,12 +7,17 @@ import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
     _DATE_COLUMNS,
     SQLAccessConfig,
     SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
     SQLUploader,
     SQLUploaderConfig,
     SQLUploadStager,
@@ -57,6 +62,57 @@ class PostgresConnectionConfig(SQLConnectionConfig):
         )
 
 
+class PostgresIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class PostgresIndexer(SQLIndexer):
+    connection_config: PostgresConnectionConfig
+    index_config: PostgresIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def _get_doc_ids(self) -> list[str]:
+        connection = self.connection_config.get_connection()
+        with connection.cursor() as cursor:
+            cursor.execute(
+                f"SELECT {self.index_config.id_column} FROM {self.index_config.table_name}"
+            )
+            results = cursor.fetchall()
+            ids = [result[0] for result in results]
+            return ids
+
+
+class PostgresDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class PostgresDownloader(SQLDownloader):
+    connection_config: PostgresConnectionConfig
+    download_config: PostgresDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
+        table_name = file_data.additional_metadata["table_name"]
+        id_column = file_data.additional_metadata["id_column"]
+        ids = file_data.additional_metadata["ids"]
+        connection = self.connection_config.get_connection()
+        with connection.cursor() as cursor:
+            fields = ",".join(self.download_config.fields) if self.download_config.fields else "*"
+            query = "SELECT {fields} FROM {table_name} WHERE {id_column} in ({ids})".format(
+                fields=fields,
+                table_name=table_name,
+                id_column=id_column,
+                ids=",".join([str(i) for i in ids]),
+            )
+            logger.debug(f"running query: {query}")
+            cursor.execute(query)
+            rows = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+            return rows, columns
+
+
 class PostgresUploadStagerConfig(SQLUploadStagerConfig):
     pass