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

test/integration/chunkers/test_chunkers.py

@@ -0,0 +1,42 @@
+import os
+from pathlib import Path
+
+import pytest
+
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.processes.chunker import Chunker, ChunkerConfig
+
+int_test_dir = Path(__file__).parent
+assets_dir = int_test_dir / "assets"
+
+chunker_files = [path for path in assets_dir.iterdir() if path.is_file()]
+
+
+@pytest.mark.parametrize("chunker_file", chunker_files, ids=[path.name for path in chunker_files])
+@pytest.mark.parametrize("strategy", ["basic", "by_title", "by_similarity", "by_page"])
+@requires_env("UNSTRUCTURED_API_KEY", "UNSTRUCTURED_API_URL")
+@pytest.mark.asyncio
+async def test_chunker_api(chunker_file: Path, strategy: str):
+    api_key = os.getenv("UNSTRUCTURED_API_KEY")
+    api_url = os.getenv("UNSTRUCTURED_API_URL")
+
+    chunker_config = ChunkerConfig(
+        chunking_strategy=strategy,
+        chunk_by_api=True,
+        chunk_api_key=api_key,
+        chunking_endpoint=api_url,
+    )
+    chunker = Chunker(config=chunker_config)
+    results = await chunker.run_async(elements_filepath=chunker_file)
+    assert results
+
+
+@pytest.mark.parametrize("chunker_file", chunker_files, ids=[path.name for path in chunker_files])
+@pytest.mark.parametrize("strategy", ["basic", "by_title"])
+def test_chunker_basic(chunker_file: Path, strategy: str):
+    chunker_config = ChunkerConfig(
+        chunking_strategy=strategy,
+    )
+    chunker = Chunker(config=chunker_config)
+    results = chunker.run(elements_filepath=chunker_file)
+    assert results

test/integration/connectors/conftest.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+
+import pytest
+
+FILENAME = "DA-1p-with-duplicate-pages.pdf.json"
+
+
+@pytest.fixture
+def upload_file() -> Path:
+    int_test_dir = Path(__file__).parent
+    assets_dir = int_test_dir / "assets"
+    upload_file = assets_dir / FILENAME
+    assert upload_file.exists()
+    assert upload_file.is_file()
+    return upload_file

test/integration/connectors/databricks_tests/test_volumes_native.py

@@ -0,0 +1,165 @@
+import json
+import os
+import tempfile
+import uuid
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+
+import pytest
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors.platform import NotFound
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG
+from test.integration.connectors.utils.validation import (
+    ValidationConfigs,
+    source_connector_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.processes.connectors.databricks.volumes_native import (
+    CONNECTOR_TYPE,
+    DatabricksNativeVolumesAccessConfig,
+    DatabricksNativeVolumesConnectionConfig,
+    DatabricksNativeVolumesDownloader,
+    DatabricksNativeVolumesDownloaderConfig,
+    DatabricksNativeVolumesIndexer,
+    DatabricksNativeVolumesIndexerConfig,
+    DatabricksNativeVolumesUploader,
+    DatabricksNativeVolumesUploaderConfig,
+)
+
+
+@dataclass
+class EnvData:
+    host: str
+    client_id: str
+    client_secret: str
+    catalog: str
+
+    def get_connection_config(self) -> DatabricksNativeVolumesConnectionConfig:
+        return DatabricksNativeVolumesConnectionConfig(
+            host=self.host,
+            access_config=DatabricksNativeVolumesAccessConfig(
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+            ),
+        )
+
+
+def get_env_data() -> EnvData:
+    return EnvData(
+        host=os.environ["DATABRICKS_HOST"],
+        client_id=os.environ["DATABRICKS_CLIENT_ID"],
+        client_secret=os.environ["DATABRICKS_CLIENT_SECRET"],
+        catalog=os.environ["DATABRICKS_CATALOG"],
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env(
+    "DATABRICKS_HOST", "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET", "DATABRICKS_CATALOG"
+)
+async def test_volumes_native_source():
+    env_data = get_env_data()
+    indexer_config = DatabricksNativeVolumesIndexerConfig(
+        recursive=True,
+        volume="test-platform",
+        volume_path="databricks-volumes-test-input",
+        catalog=env_data.catalog,
+    )
+    connection_config = env_data.get_connection_config()
+    with tempfile.TemporaryDirectory() as tempdir:
+        tempdir_path = Path(tempdir)
+        download_config = DatabricksNativeVolumesDownloaderConfig(download_dir=tempdir_path)
+        indexer = DatabricksNativeVolumesIndexer(
+            connection_config=connection_config, index_config=indexer_config
+        )
+        downloader = DatabricksNativeVolumesDownloader(
+            connection_config=connection_config, download_config=download_config
+        )
+        await source_connector_validation(
+            indexer=indexer,
+            downloader=downloader,
+            configs=ValidationConfigs(
+                test_id="databricks_volumes_native",
+                expected_num_files=1,
+            ),
+        )
+
+
+def _get_volume_path(catalog: str, volume: str, volume_path: str):
+    return f"/Volumes/{catalog}/default/{volume}/{volume_path}"
+
+
+@contextmanager
+def databricks_destination_context(env_data: EnvData, volume: str, volume_path) -> WorkspaceClient:
+    client = WorkspaceClient(
+        host=env_data.host, client_id=env_data.client_id, client_secret=env_data.client_secret
+    )
+    try:
+        yield client
+    finally:
+        # Cleanup
+        try:
+            for file in client.files.list_directory_contents(
+                directory_path=_get_volume_path(env_data.catalog, volume, volume_path)
+            ):
+                client.files.delete(file.path)
+            client.files.delete_directory(_get_volume_path(env_data.catalog, volume, volume_path))
+        except NotFound:
+            # Directory was never created, don't need to delete
+            pass
+
+
+def validate_upload(client: WorkspaceClient, catalog: str, volume: str, volume_path: str):
+    files = list(
+        client.files.list_directory_contents(
+            directory_path=_get_volume_path(catalog, volume, volume_path)
+        )
+    )
+
+    assert len(files) == 1
+
+    resp = client.files.download(files[0].path)
+    data = json.loads(resp.contents.read())
+
+    assert len(data) == 22
+    element_types = {v["type"] for v in data}
+    assert len(element_types) == 1
+    assert "CompositeElement" in element_types
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env(
+    "DATABRICKS_HOST", "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET", "DATABRICKS_CATALOG"
+)
+async def test_volumes_native_destination(upload_file: Path):
+    env_data = get_env_data()
+    volume_path = f"databricks-volumes-test-output-{uuid.uuid4()}"
+    mock_file_data = FileData(identifier="mock file data", connector_type=CONNECTOR_TYPE)
+    with databricks_destination_context(
+        volume="test-platform", volume_path=volume_path, env_data=env_data
+    ) as workspace_client:
+        connection_config = env_data.get_connection_config()
+        uploader = DatabricksNativeVolumesUploader(
+            connection_config=connection_config,
+            upload_config=DatabricksNativeVolumesUploaderConfig(
+                volume="test-platform",
+                volume_path=volume_path,
+                catalog=env_data.catalog,
+            ),
+        )
+        if uploader.is_async():
+            await uploader.run_async(path=upload_file, file_data=mock_file_data)
+        else:
+            uploader.run(path=upload_file, file_data=mock_file_data)
+
+        validate_upload(
+            client=workspace_client,
+            catalog=env_data.catalog,
+            volume="test-platform",
+            volume_path=volume_path,
+        )

test/integration/connectors/sql/test_postgres.py

@@ -0,0 +1,178 @@
+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, 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,
+    expected_num_elements: int,
+    test_embedding: list[float],
+    expected_text: str,
+):
+    # Run the following validations:
+    # * Check that the number of records in the table match the expected value
+    # * Given the embedding, make sure it matches the associated text it belongs to
+    with connect(**connect_params) as connection:
+        cursor = connection.cursor()
+        query = "select count(*) from elements;"
+        cursor.execute(query)
+        count = cursor.fetchone()[0]
+        assert (
+            count == expected_num_elements
+        ), f"dest check failed: got {count}, expected {expected_num_elements}"
+
+        cursor.execute("SELECT embeddings FROM elements order by text limit 1")
+        similarity_query = (
+            f"SELECT text FROM elements ORDER BY embeddings <-> '{test_embedding}' LIMIT 1;"
+        )
+        cursor.execute(similarity_query)
+        res = cursor.fetchone()
+        assert res[0] == expected_text
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "sql")
+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" / "postgres" / "destination"
+    ):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            stager = PostgresUploadStager()
+            stager_params = {
+                "elements_filepath": upload_file,
+                "file_data": mock_file_data,
+                "output_dir": Path(tmpdir),
+                "output_filename": "test_db",
+            }
+            if stager.is_async():
+                staged_path = await stager.run_async(**stager_params)
+            else:
+                staged_path = stager.run(**stager_params)
+
+            # The stager should append the `.json` suffix to the output filename passed in.
+            assert staged_path.name == "test_db.json"
+
+            connect_params = {
+                "host": "localhost",
+                "port": 5433,
+                "database": "elements",
+                "user": "unstructured",
+                "password": "test",
+            }
+
+            uploader = PostgresUploader(
+                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"]),
+                )
+            )
+            if uploader.is_async():
+                await uploader.run_async(path=staged_path, file_data=mock_file_data)
+            else:
+                uploader.run(path=staged_path, file_data=mock_file_data)
+
+            staged_df = pd.read_json(staged_path, orient="records", lines=True)
+            sample_element = staged_df.iloc[0]
+            expected_num_elements = len(staged_df)
+            validate_destination(
+                connect_params=connect_params,
+                expected_num_elements=expected_num_elements,
+                expected_text=sample_element["text"],
+                test_embedding=sample_element["embeddings"],
+            )

test/integration/connectors/sql/test_sqlite.py

@@ -0,0 +1,151 @@
+import sqlite3
+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, 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_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" / "destination" / "sqlite-schema.sql"
+        assert db_init_path.exists()
+        assert db_init_path.is_file()
+        with sqlite3.connect(database=db_path) as sqlite_connection:
+            with db_init_path.open("r") as f:
+                query = f.read()
+            cursor = sqlite_connection.cursor()
+            cursor.executescript(query)
+        yield db_path
+
+
+def validate_destination(db_path: Path, expected_num_elements: int):
+    # Run the following validations:
+    # * Check that the number of records in the table match the expected value
+    connection = None
+    try:
+        connection = sqlite3.connect(database=db_path)
+        query = "select count(*) from elements;"
+        cursor = connection.cursor()
+        cursor.execute(query)
+        count = cursor.fetchone()[0]
+        assert (
+            count == expected_num_elements
+        ), f"dest check failed: got {count}, expected {expected_num_elements}"
+    finally:
+        if connection:
+            connection.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "sql")
+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_upload_setup() as db_path:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            stager = SQLiteUploadStager()
+            stager_params = {
+                "elements_filepath": upload_file,
+                "file_data": mock_file_data,
+                "output_dir": Path(tmpdir),
+                "output_filename": "test_db",
+            }
+            if stager.is_async():
+                staged_path = await stager.run_async(**stager_params)
+            else:
+                staged_path = stager.run(**stager_params)
+
+            # The stager should append the `.json` suffix to the output filename passed in.
+            assert staged_path.name == "test_db.json"
+
+            uploader = SQLiteUploader(
+                connection_config=SQLiteConnectionConfig(database_path=db_path)
+            )
+            if uploader.is_async():
+                await uploader.run_async(path=staged_path, file_data=mock_file_data)
+            else:
+                uploader.run(path=staged_path, file_data=mock_file_data)
+
+            staged_df = pd.read_json(staged_path, orient="records", lines=True)
+            validate_destination(db_path=db_path, expected_num_elements=len(staged_df))