unstructured-ingest 0.0.24__py3-none-any.whl → 0.1.0__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/test_postgres.py

@@ -0,0 +1,100 @@
+import tempfile
+from pathlib import Path
+
+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.docker_compose import docker_compose_context
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.processes.connectors.sql.postgres import (
+    CONNECTOR_TYPE,
+    PostgresAccessConfig,
+    PostgresConnectionConfig,
+    PostgresUploader,
+    PostgresUploadStager,
+)
+
+
+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"):
+        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/test_s3.py

@@ -0,0 +1,152 @@
+import os
+import tempfile
+import uuid
+from pathlib import Path
+
+import pytest
+
+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 test.integration.utils import requires_env
+from unstructured_ingest.error import (
+    SourceConnectionError,
+)
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.fsspec.s3 import (
+    CONNECTOR_TYPE,
+    S3AccessConfig,
+    S3ConnectionConfig,
+    S3Downloader,
+    S3DownloaderConfig,
+    S3Indexer,
+    S3IndexerConfig,
+    S3Uploader,
+    S3UploaderConfig,
+)
+
+
+def validate_predownload_file_data(file_data: FileData):
+    assert file_data.connector_type == CONNECTOR_TYPE
+    assert file_data.local_download_path is None
+
+
+def validate_postdownload_file_data(file_data: FileData):
+    assert file_data.connector_type == CONNECTOR_TYPE
+    assert file_data.local_download_path is not None
+
+
+@pytest.fixture
+def anon_connection_config() -> S3ConnectionConfig:
+    return S3ConnectionConfig(access_config=S3AccessConfig(), anonymous=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+async def test_s3_source(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-dev-tech-fixtures/small-pdf-set/")
+    with tempfile.TemporaryDirectory() as tempdir:
+        tempdir_path = Path(tempdir)
+        download_config = S3DownloaderConfig(download_dir=tempdir_path)
+        indexer = S3Indexer(connection_config=anon_connection_config, index_config=indexer_config)
+        downloader = S3Downloader(
+            connection_config=anon_connection_config, download_config=download_config
+        )
+        await source_connector_validation(
+            indexer=indexer,
+            downloader=downloader,
+            configs=ValidationConfigs(
+                test_id="s3",
+                predownload_file_data_check=validate_predownload_file_data,
+                postdownload_file_data_check=validate_postdownload_file_data,
+                expected_num_files=4,
+            ),
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+async def test_s3_source_no_access(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-ingest-test-fixtures/destination/")
+    indexer = S3Indexer(connection_config=anon_connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "minio")
+async def test_s3_minio_source(anon_connection_config: S3ConnectionConfig):
+    anon_connection_config.endpoint_url = "http://localhost:9000"
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-dev-tech-fixtures/")
+    with docker_compose_context(docker_compose_path=env_setup_path / "minio"):
+        with tempfile.TemporaryDirectory() as tempdir:
+            tempdir_path = Path(tempdir)
+            download_config = S3DownloaderConfig(download_dir=tempdir_path)
+            indexer = S3Indexer(
+                connection_config=anon_connection_config, index_config=indexer_config
+            )
+            downloader = S3Downloader(
+                connection_config=anon_connection_config, download_config=download_config
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="s3-minio",
+                    predownload_file_data_check=validate_predownload_file_data,
+                    postdownload_file_data_check=validate_postdownload_file_data,
+                    expected_num_files=1,
+                    exclude_fields_extend=[
+                        "metadata.date_modified",
+                        "metadata.date_created",
+                        "additional_metadata.LastModified",
+                    ],
+                ),
+            )
+
+
+def get_aws_credentials() -> dict:
+    access_key = os.getenv("S3_INGEST_TEST_ACCESS_KEY", None)
+    assert access_key
+    secret_key = os.getenv("S3_INGEST_TEST_SECRET_KEY", None)
+    assert secret_key
+    return {"aws_access_key_id": access_key, "aws_secret_access_key": secret_key}
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("S3_INGEST_TEST_ACCESS_KEY", "S3_INGEST_TEST_SECRET_KEY")
+async def test_s3_destination(upload_file: Path):
+    aws_credentials = get_aws_credentials()
+    s3_bucket = "s3://utic-ingest-test-fixtures"
+    destination_path = f"{s3_bucket}/destination/{uuid.uuid4()}"
+    connection_config = S3ConnectionConfig(
+        access_config=S3AccessConfig(
+            key=aws_credentials["aws_access_key_id"],
+            secret=aws_credentials["aws_secret_access_key"],
+        ),
+    )
+    upload_config = S3UploaderConfig(remote_url=destination_path)
+    uploader = S3Uploader(connection_config=connection_config, upload_config=upload_config)
+    s3fs = uploader.fs
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    try:
+        if uploader.is_async():
+            await uploader.run_async(path=upload_file, file_data=file_data)
+        else:
+            uploader.run(path=upload_file, file_data=file_data)
+        uploaded_files = s3fs.ls(path=destination_path)
+        assert len(uploaded_files) == 1
+    finally:
+        s3fs.rm(path=destination_path, recursive=True)

test/integration/connectors/test_sqlite.py

@@ -0,0 +1,91 @@
+import sqlite3
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.processes.connectors.sql.sqlite import (
+    CONNECTOR_TYPE,
+    SQLiteConnectionConfig,
+    SQLiteUploader,
+    SQLiteUploadStager,
+)
+
+
+@contextmanager
+def sqlite_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"
+        assert db_init_path.exists()
+        assert db_init_path.is_file()
+        connection = None
+        try:
+            connection = sqlite3.connect(database=db_path)
+            with db_init_path.open("r") as f:
+                query = f.read()
+            cursor = connection.cursor()
+            cursor.executescript(query)
+            yield db_path
+        finally:
+            if connection:
+                connection.close()
+
+
+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_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))

test/integration/connectors/utils/constants.py

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

test/integration/connectors/utils/docker_compose.py

@@ -0,0 +1,44 @@
+import subprocess
+from contextlib import contextmanager
+from pathlib import Path
+
+
+@contextmanager
+def docker_compose_context(docker_compose_path: Path):
+    # Dynamically run a specific docker compose file and make sure it gets cleanup by
+    # by leveraging a context manager. Uses subprocess to map docker compose commands
+    # to the underlying shell.
+    assert docker_compose_path.exists()
+    if docker_compose_path.is_dir():
+        if (docker_compose_path / "docker-compose.yml").exists():
+            docker_compose_path = docker_compose_path / "docker-compose.yml"
+        elif (docker_compose_path / "docker-compose.yaml").exists():
+            docker_compose_path = docker_compose_path / "docker-compose.yaml"
+    assert docker_compose_path.is_file()
+    resp = None
+    try:
+        cmd = f"docker compose -f {docker_compose_path.resolve()} up -d --wait"
+        print(f"Running command: {cmd}")
+        resp = subprocess.run(
+            cmd,
+            shell=True,
+            capture_output=True,
+        )
+        # Return code from docker compose using --wait can be 1 even if no error
+        yield
+    except Exception as e:
+        if resp:
+            print("STDOUT: {}".format(resp.stdout.decode("utf-8")))
+            print("STDERR: {}".format(resp.stderr.decode("utf-8")))
+        raise e
+    finally:
+        cmd = f"docker compose -f {docker_compose_path.resolve()} down --remove-orphans -v"
+        print(f"Running command: {cmd}")
+        final_resp = subprocess.run(
+            cmd,
+            shell=True,
+            capture_output=True,
+        )
+        if final_resp.returncode != 0:
+            print("STDOUT: {}".format(final_resp.stdout.decode("utf-8")))
+            print("STDERR: {}".format(final_resp.stderr.decode("utf-8")))