unstructured-ingest 0.3.13__py3-none-any.whl

test/integration/chunkers/test_chunkers.py

@@ -0,0 +1,31 @@
+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

test/integration/connectors/conftest.py

@@ -0,0 +1,38 @@
+import tempfile
+from pathlib import Path
+from typing import Generator
+
+import pytest
+
+from unstructured_ingest.v2.logger import logger
+
+FILENAME = Path("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
+
+
+@pytest.fixture
+def upload_file_ndjson() -> Path:
+    int_test_dir = Path(__file__).parent
+    assets_dir = int_test_dir / "assets"
+    upload_file = assets_dir / FILENAME.with_suffix(".ndjson")
+    assert upload_file.exists()
+    assert upload_file.is_file()
+    return upload_file
+
+
+@pytest.fixture
+def temp_dir() -> Generator[Path, None, None]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+        logger.info(f"Created temp dir '{temp_path}'")
+        yield temp_path
+        logger.info(f"Removing temp dir '{temp_path}'")

test/integration/connectors/databricks/test_volumes_native.py

@@ -0,0 +1,269 @@
+import json
+import os
+import uuid
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from unittest import mock
+
+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.source import (
+    SourceValidationConfigs,
+    source_connector_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.errors import UserAuthError, UserError
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.databricks.volumes_native import (
+    CONNECTOR_TYPE,
+    DatabricksNativeVolumesAccessConfig,
+    DatabricksNativeVolumesConnectionConfig,
+    DatabricksNativeVolumesDownloader,
+    DatabricksNativeVolumesDownloaderConfig,
+    DatabricksNativeVolumesIndexer,
+    DatabricksNativeVolumesIndexerConfig,
+    DatabricksNativeVolumesUploader,
+    DatabricksNativeVolumesUploaderConfig,
+)
+
+
+@dataclass
+class BaseEnvData:
+    host: str
+    catalog: str
+
+
+@dataclass
+class BasicAuthEnvData(BaseEnvData):
+    client_id: str
+    client_secret: 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,
+            ),
+        )
+
+
+@dataclass
+class PATEnvData(BaseEnvData):
+    token: str
+
+    def get_connection_config(self) -> DatabricksNativeVolumesConnectionConfig:
+        return DatabricksNativeVolumesConnectionConfig(
+            host=self.host,
+            access_config=DatabricksNativeVolumesAccessConfig(
+                token=self.token,
+            ),
+        )
+
+
+def get_basic_auth_env_data() -> BasicAuthEnvData:
+    return BasicAuthEnvData(
+        host=os.environ["DATABRICKS_HOST"],
+        client_id=os.environ["DATABRICKS_CLIENT_ID"],
+        client_secret=os.environ["DATABRICKS_CLIENT_SECRET"],
+        catalog=os.environ["DATABRICKS_CATALOG"],
+    )
+
+
+def get_pat_env_data() -> PATEnvData:
+    return PATEnvData(
+        host=os.environ["DATABRICKS_HOST"],
+        catalog=os.environ["DATABRICKS_CATALOG"],
+        token=os.environ["DATABRICKS_PAT"],
+    )
+
+
+@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(tmp_path: Path):
+    env_data = get_basic_auth_env_data()
+    with mock.patch.dict(os.environ, clear=True):
+        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()
+        download_config = DatabricksNativeVolumesDownloaderConfig(download_dir=tmp_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=SourceValidationConfigs(
+                test_id="databricks_volumes_native",
+                expected_num_files=1,
+            ),
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("DATABRICKS_HOST", "DATABRICKS_PAT", "DATABRICKS_CATALOG")
+async def test_volumes_native_source_pat(tmp_path: Path):
+    env_data = get_pat_env_data()
+    with mock.patch.dict(os.environ, clear=True):
+        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()
+        download_config = DatabricksNativeVolumesDownloaderConfig(download_dir=tmp_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=SourceValidationConfigs(
+                test_id="databricks_volumes_native_pat",
+                expected_num_files=1,
+            ),
+        )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("DATABRICKS_HOST", "DATABRICKS_PAT", "DATABRICKS_CATALOG")
+def test_volumes_native_source_pat_invalid_catalog():
+    env_data = get_pat_env_data()
+    with mock.patch.dict(os.environ, clear=True):
+        indexer_config = DatabricksNativeVolumesIndexerConfig(
+            recursive=True,
+            volume="test-platform",
+            volume_path="databricks-volumes-test-input",
+            catalog="fake_catalog",
+        )
+        indexer = DatabricksNativeVolumesIndexer(
+            connection_config=env_data.get_connection_config(), index_config=indexer_config
+        )
+        with pytest.raises(UserError):
+            _ = list(indexer.run())
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("DATABRICKS_HOST")
+def test_volumes_native_source_pat_invalid_pat():
+    host = os.environ["DATABRICKS_HOST"]
+    with mock.patch.dict(os.environ, clear=True):
+        indexer_config = DatabricksNativeVolumesIndexerConfig(
+            recursive=True,
+            volume="test-platform",
+            volume_path="databricks-volumes-test-input",
+            catalog="fake_catalog",
+        )
+        connection_config = DatabricksNativeVolumesConnectionConfig(
+            host=host,
+            access_config=DatabricksNativeVolumesAccessConfig(
+                token="invalid-token",
+            ),
+        )
+        indexer = DatabricksNativeVolumesIndexer(
+            connection_config=connection_config, index_config=indexer_config
+        )
+        with pytest.raises(UserAuthError):
+            _ = list(indexer.run())
+
+
+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: BasicAuthEnvData, 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_basic_auth_env_data()
+    volume_path = f"databricks-volumes-test-output-{uuid.uuid4()}"
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    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,
+            ),
+        )
+        uploader.precheck()
+        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)
+
+        validate_upload(
+            client=workspace_client,
+            catalog=env_data.catalog,
+            volume="test-platform",
+            volume_path=volume_path,
+        )

test/integration/connectors/discord/test_discord.py

@@ -0,0 +1,90 @@
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import pytest
+
+from test.integration.connectors.utils.constants import SOURCE_TAG
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
+    source_connector_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.error import SourceConnectionError
+from unstructured_ingest.v2.processes.connectors.discord import (
+    CONNECTOR_TYPE,
+    DiscordAccessConfig,
+    DiscordConnectionConfig,
+    DiscordDownloader,
+    DiscordDownloaderConfig,
+    DiscordIndexer,
+    DiscordIndexerConfig,
+)
+
+
+@dataclass(frozen=True)
+class EnvData:
+    token: Optional[str]
+    channels: Optional[list[str]]
+
+
+def get_env_data() -> EnvData:
+    return EnvData(
+        token=os.getenv("DISCORD_TOKEN"),
+        channels=os.getenv("DISCORD_CHANNELS", default=[]).split(","),
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("DISCORD_TOKEN", "DISCORD_CHANNELS")
+async def test_discord_source():
+    env = get_env_data()
+    indexer_config = DiscordIndexerConfig(channels=env.channels)
+    with tempfile.TemporaryDirectory() as tempdir:
+        tempdir_path = Path(tempdir)
+        connection_config = DiscordConnectionConfig(
+            access_config=DiscordAccessConfig(token=env.token)
+        )
+        download_config = DiscordDownloaderConfig(download_dir=tempdir_path)
+        indexer = DiscordIndexer(connection_config=connection_config, index_config=indexer_config)
+        downloader = DiscordDownloader(
+            connection_config=connection_config, download_config=download_config
+        )
+        expected_num_files = len(env.channels)
+        await source_connector_validation(
+            indexer=indexer,
+            downloader=downloader,
+            configs=SourceValidationConfigs(
+                test_id=CONNECTOR_TYPE,
+                expected_num_files=expected_num_files,
+                expected_number_indexed_file_data=expected_num_files,
+                validate_downloaded_files=True,
+            ),
+        )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("DISCORD_CHANNELS")
+def test_discord_source_precheck_fail_no_token():
+    indexer_config = DiscordIndexerConfig(channels=get_env_data().channels)
+
+    connection_config = DiscordConnectionConfig(access_config=DiscordAccessConfig(token=""))
+    indexer = DiscordIndexer(connection_config=connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("DISCORD_TOKEN")
+def test_discord_source_precheck_fail_no_channels():
+    indexer_config = DiscordIndexerConfig(channels=[])
+
+    connection_config = DiscordConnectionConfig(
+        access_config=DiscordAccessConfig(token=get_env_data().token)
+    )
+    indexer = DiscordIndexer(connection_config=connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()

test/integration/connectors/duckdb/conftest.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+
+import pytest
+
+int_test_dir = Path(__file__).parent
+assets_dir = int_test_dir / "assets"
+
+
+@pytest.fixture
+def duckdb_schema() -> Path:
+    schema_file = assets_dir / "duckdb-schema.sql"
+    assert schema_file.exists()
+    assert schema_file.is_file()
+    return schema_file

test/integration/connectors/duckdb/test_duckdb.py

@@ -0,0 +1,89 @@
+import json
+from pathlib import Path
+
+import duckdb
+import pytest
+from _pytest.fixtures import TopRequest
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from test.integration.connectors.utils.validation.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.duckdb.duckdb import (
+    CONNECTOR_TYPE,
+    DuckDBConnectionConfig,
+    DuckDBUploader,
+    DuckDBUploaderConfig,
+    DuckDBUploadStager,
+)
+
+
+@pytest.fixture
+def provisioned_db_file(duckdb_schema: Path, temp_dir: Path) -> Path:
+    db_path = Path(temp_dir) / "temp_duck.db"
+    with duckdb.connect(database=db_path) as duckdb_connection:
+        with duckdb_schema.open("r") as f:
+            query = f.read()
+        duckdb_connection.execute(query)
+        duckdb_connection.close()
+    return db_path
+
+
+def validate_duckdb_destination(db_path: Path, expected_num_elements: int):
+    conn = None
+    try:
+        conn = duckdb.connect(db_path)
+        _results = conn.sql("select count(*) from elements").fetchall()
+        _count = _results[0][0]
+        assert (
+            _count == expected_num_elements
+        ), f"dest check failed: got {_count}, expected {expected_num_elements}"
+        conn.close()
+    finally:
+        if conn:
+            conn.close()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "duckdb")
+def test_duckdb_destination(upload_file: Path, provisioned_db_file: Path, temp_dir: Path):
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    stager = DuckDBUploadStager()
+    staged_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    connection_config = DuckDBConnectionConfig(database=str(provisioned_db_file))
+    upload_config = DuckDBUploaderConfig()
+    uploader = DuckDBUploader(connection_config=connection_config, upload_config=upload_config)
+
+    uploader.run(path=staged_path, file_data=file_data)
+
+    with staged_path.open() as f:
+        data = json.load(f)
+    validate_duckdb_destination(db_path=provisioned_db_file, expected_num_elements=len(data))
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_duckdb_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = DuckDBUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/duckdb/test_motherduck.py

@@ -0,0 +1,95 @@
+import os
+import uuid
+from pathlib import Path
+from typing import Generator
+
+import duckdb
+import pandas as pd
+import pytest
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.duckdb.motherduck import (
+    CONNECTOR_TYPE,
+    MotherDuckAccessConfig,
+    MotherDuckConnectionConfig,
+    MotherDuckUploader,
+    MotherDuckUploaderConfig,
+    MotherDuckUploadStager,
+)
+
+
+@pytest.fixture
+def md_token() -> str:
+    motherduck_token = os.getenv("MOTHERDUCK_TOKEN", None)
+    assert motherduck_token
+    return motherduck_token
+
+
+@pytest.fixture
+def provisioned_db(md_token: str, duckdb_schema: Path) -> Generator[str, None, None]:
+    database_name = f"test_{str(uuid.uuid4()).replace('-', '_')}"
+    try:
+        with duckdb.connect(f"md:?motherduck_token={md_token}") as md_conn:
+            with duckdb_schema.open("r") as f:
+                query = f.read()
+            md_conn.execute(f"CREATE DATABASE {database_name}")
+            md_conn.execute(f"USE {database_name}")
+            md_conn.execute(query)
+            md_conn.close()
+        yield database_name
+    finally:
+        with duckdb.connect(f"md:?motherduck_token={md_token}") as md_conn:
+            md_conn.execute(f"DROP DATABASE {database_name}")
+            md_conn.close()
+
+
+def validate_motherduck_destination(database: str, expected_num_elements: int, md_token: str):
+    conn = None
+    try:
+        conn = duckdb.connect(f"md:?motherduck_token={md_token}")
+        conn.execute(f"USE {database}")
+        _results = conn.sql("select count(*) from elements").fetchall()
+        _count = _results[0][0]
+        assert (
+            _count == expected_num_elements
+        ), f"dest check failed: got {_count}, expected {expected_num_elements}"
+        conn.close()
+    finally:
+        if conn:
+            conn.close()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("MOTHERDUCK_TOKEN")
+def test_motherduck_destination(
+    md_token: str, upload_file: Path, provisioned_db: str, temp_dir: Path
+):
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    stager = MotherDuckUploadStager()
+    staged_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    access_config = MotherDuckAccessConfig(md_token=md_token)
+    connection_config = MotherDuckConnectionConfig(
+        database=provisioned_db, access_config=access_config
+    )
+    upload_config = MotherDuckUploaderConfig()
+    uploader = MotherDuckUploader(connection_config=connection_config, upload_config=upload_config)
+
+    uploader.run(path=staged_path, file_data=file_data)
+
+    staged_df = pd.read_json(staged_path, orient="records", lines=True)
+    validate_motherduck_destination(
+        database=provisioned_db, expected_num_elements=len(staged_df), md_token=md_token
+    )

test/integration/connectors/elasticsearch/conftest.py

@@ -0,0 +1,34 @@
+import json
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+int_test_dir = Path(__file__).parent
+assets_dir = int_test_dir / "assets"
+
+
+@pytest.fixture
+def movies_dataframe() -> pd.DataFrame:
+    movies_file = assets_dir / "wiki_movie_plots_small.csv"
+    assert movies_file.exists()
+    assert movies_file.is_file()
+    return pd.read_csv(movies_file).dropna().reset_index()
+
+
+@pytest.fixture
+def opensearch_elements_mapping() -> dict:
+    elements_mapping_file = assets_dir / "opensearch_elements_mappings.json"
+    assert elements_mapping_file.exists()
+    assert elements_mapping_file.is_file()
+    with elements_mapping_file.open() as fp:
+        return json.load(fp)
+
+
+@pytest.fixture
+def elasticsearch_elements_mapping() -> dict:
+    elements_mapping_file = assets_dir / "elasticsearch_elements_mappings.json"
+    assert elements_mapping_file.exists()
+    assert elements_mapping_file.is_file()
+    with elements_mapping_file.open() as fp:
+        return json.load(fp)