unstructured-ingest 0.7.1__py3-none-any.whl → 1.0.1__py3-none-any.whl

test/integration/connectors/databricks/test_volumes_native.py

@@ -1,273 +0,0 @@
-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 (
-    BLOB_STORAGE_TAG,
-    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.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.errors_v2 import UserAuthError, UserError
-from unstructured_ingest.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, BLOB_STORAGE_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, BLOB_STORAGE_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, BLOB_STORAGE_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, BLOB_STORAGE_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, BLOB_STORAGE_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

@@ -1,90 +0,0 @@
-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, UNCATEGORIZED_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.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, UNCATEGORIZED_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, UNCATEGORIZED_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, UNCATEGORIZED_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

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

@@ -1,90 +0,0 @@
-import json
-from pathlib import Path
-
-import duckdb
-import pytest
-from _pytest.fixtures import TopRequest
-
-from test.integration.connectors.utils.constants import DESTINATION_TAG, SQL_TAG
-from test.integration.connectors.utils.validation.destination import (
-    StagerValidationConfigs,
-    stager_validation,
-)
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.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", SQL_TAG)
-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.tags(CONNECTOR_TYPE, DESTINATION_TAG, "duckdb", SQL_TAG)
-@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

@@ -1,95 +0,0 @@
-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, SQL_TAG
-from test.integration.utils import requires_env
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.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, SQL_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

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