unstructured-ingest 0.3.13__py3-none-any.whl

test/integration/connectors/test_s3.py

@@ -0,0 +1,183 @@
+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.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.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=SourceValidationConfigs(
+                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_special_char(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://utic-dev-tech-fixtures/special-characters/")
+    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=SourceValidationConfigs(
+                test_id="s3-specialchar",
+                predownload_file_data_check=validate_predownload_file_data,
+                postdownload_file_data_check=validate_postdownload_file_data,
+                expected_num_files=1,
+            ),
+        )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+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(UserAuthError):
+        indexer.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+def test_s3_source_no_bucket(anon_connection_config: S3ConnectionConfig):
+    indexer_config = S3IndexerConfig(remote_url="s3://fake-bucket")
+    indexer = S3Indexer(connection_config=anon_connection_config, index_config=indexer_config)
+    with pytest.raises(UserError):
+        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" / "source"):
+        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=SourceValidationConfigs(
+                    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:
+        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)
+        uploaded_files = [
+            Path(file) for file in s3fs.ls(path=destination_path) if Path(file).name != "_empty"
+        ]
+        assert len(uploaded_files) == 1
+    finally:
+        s3fs.rm(path=destination_path, recursive=True)

test/integration/connectors/test_vectara.py

@@ -0,0 +1,270 @@
+import json
+import os
+import time
+from pathlib import Path
+from typing import Generator
+from uuid import uuid4
+
+import pytest
+import requests
+
+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.logger import logger
+from unstructured_ingest.v2.processes.connectors.vectara import (
+    CONNECTOR_TYPE as VECTARA_CONNECTOR_TYPE,
+)
+from unstructured_ingest.v2.processes.connectors.vectara import (
+    VectaraAccessConfig,
+    VectaraConnectionConfig,
+    VectaraUploader,
+    VectaraUploaderConfig,
+    VectaraUploadStager,
+    VectaraUploadStagerConfig,
+)
+
+
+def validate_upload(response: dict, expected_data: dict):
+    element_id = expected_data["element_id"]
+    expected_text = expected_data["text"]
+    filename = expected_data["metadata"]["filename"]
+    filetype = expected_data["metadata"]["filetype"]
+    page_number = expected_data["metadata"]["page_number"]
+
+    response = response["search_results"][0]
+
+    assert response is not None
+    assert response["text"] == expected_text
+    assert response["part_metadata"]["element_id"] == element_id
+    assert response["part_metadata"]["filename"] == filename
+    assert response["part_metadata"]["filetype"] == filetype
+    assert response["part_metadata"]["page_number"] == page_number
+
+
+@requires_env("VECTARA_OAUTH_CLIENT_ID", "VECTARA_OAUTH_SECRET", "VECTARA_CUSTOMER_ID")
+def _get_jwt_token():
+    """Connect to the server and get a JWT token."""
+    customer_id = os.environ["VECTARA_CUSTOMER_ID"]
+    token_endpoint = (
+        f"https://vectara-prod-{customer_id}.auth.us-west-2.amazoncognito.com/oauth2/token"
+    )
+    headers = {
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+    data = {
+        "grant_type": "client_credentials",
+        "client_id": os.environ["VECTARA_OAUTH_CLIENT_ID"],
+        "client_secret": os.environ["VECTARA_OAUTH_SECRET"],
+    }
+
+    response = requests.post(token_endpoint, headers=headers, data=data)
+    response.raise_for_status()
+    response_json = response.json()
+
+    return response_json.get("access_token")
+
+
+def query_data(corpus_key: str, element_id: str) -> dict:
+
+    url = f"https://api.vectara.io/v2/corpora/{corpus_key}/query"
+
+    # the query below requires the corpus to have filter attributes for element_id
+
+    data = json.dumps(
+        {
+            "query": "string",
+            "search": {
+                "metadata_filter": f"part.element_id = '{element_id}'",
+                "lexical_interpolation": 1,
+                "limit": 10,
+            },
+        }
+    )
+
+    jwt_token = _get_jwt_token()
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {jwt_token}",
+        "X-source": "unstructured",
+    }
+
+    response = requests.post(url, headers=headers, data=data)
+    response.raise_for_status()
+    response_json = response.json()
+
+    return response_json
+
+
+def create_corpora(corpus_key: str, corpus_name: str) -> None:
+    url = "https://api.vectara.io/v2/corpora"
+    data = json.dumps({"key": corpus_key, "name": corpus_name, "description": "integration test"})
+    jwt_token = _get_jwt_token()
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {jwt_token}",
+        "X-source": "unstructured",
+    }
+
+    response = requests.post(url, headers=headers, data=data)
+    response.raise_for_status()
+
+
+def replace_filter_attributes(corpus_key: str) -> None:
+    url = f"https://api.vectara.io/v2/corpora/{corpus_key}/replace_filter_attributes"
+    data = json.dumps(
+        {
+            "filter_attributes": [
+                {"name": "element_id", "level": "part", "indexed": True, "type": "text"}
+            ]
+        }
+    )
+    jwt_token = _get_jwt_token()
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {jwt_token}",
+        "X-source": "unstructured",
+    }
+
+    response = requests.post(url, headers=headers, data=data)
+    response.raise_for_status()
+
+
+def delete_corpora(corpus_key: str) -> None:
+    url = f"https://api.vectara.io/v2/corpora/{corpus_key}"
+
+    jwt_token = _get_jwt_token()
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {jwt_token}",
+        "X-source": "unstructured",
+    }
+
+    response = requests.delete(url, headers=headers)
+    response.raise_for_status()
+
+
+def list_corpora() -> list:
+    url = "https://api.vectara.io/v2/corpora?limit=100"
+    jwt_token = _get_jwt_token()
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {jwt_token}",
+        "X-source": "unstructured",
+    }
+    response = requests.get(url, headers=headers)
+    response.raise_for_status()
+    response_json = response.json()
+    if response_json.get("corpora"):
+        return [item["key"] for item in response_json.get("corpora")]
+    else:
+        return []
+
+
+def wait_for_ready(corpus_key: str, timeout=60, interval=2) -> None:
+    def is_ready_status():
+        corpora_list = list_corpora()
+        return corpus_key in corpora_list
+
+    start = time.time()
+    is_ready = is_ready_status()
+    while not is_ready and time.time() - start < timeout:
+        time.sleep(interval)
+        is_ready = is_ready_status()
+    if not is_ready:
+        raise TimeoutError("time out waiting for corpus to be ready")
+
+
+def wait_for_delete(corpus_key: str, timeout=60, interval=2) -> None:
+    start = time.time()
+    while time.time() - start < timeout:
+        corpora_list = list_corpora()
+        if corpus_key not in corpora_list:
+            return
+        time.sleep(interval)
+
+    raise TimeoutError("time out waiting for corpus to delete")
+
+
+@pytest.fixture
+def corpora_util() -> Generator[str, None, None]:
+    random_id = str(uuid4()).split("-")[0]
+    corpus_key = f"ingest-test-{random_id}"
+    corpus_name = "ingest-test"
+    logger.info(f"Creating corpus with key: {corpus_key}")
+    try:
+        create_corpora(corpus_key, corpus_name)
+        replace_filter_attributes(corpus_key)
+        wait_for_ready(corpus_key=corpus_key)
+        yield corpus_key
+    except Exception as e:
+        logger.error(f"failed to create corpus {corpus_key}: {e}")
+    finally:
+        logger.info(f"deleting corpus: {corpus_key}")
+        delete_corpora(corpus_key)
+        wait_for_delete(corpus_key=corpus_key)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(VECTARA_CONNECTOR_TYPE, DESTINATION_TAG, "vectara")
+@requires_env("VECTARA_OAUTH_CLIENT_ID", "VECTARA_OAUTH_SECRET", "VECTARA_CUSTOMER_ID")
+async def test_vectara_destination(
+    upload_file: Path, tmp_path: Path, corpora_util: str, retries=30, interval=10
+):
+    corpus_key = corpora_util
+    connection_kwargs = {
+        "customer_id": os.environ["VECTARA_CUSTOMER_ID"],
+        "corpus_key": corpus_key,
+    }
+
+    oauth_client_id = os.environ["VECTARA_OAUTH_CLIENT_ID"]
+    oauth_secret = os.environ["VECTARA_OAUTH_SECRET"]
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=VECTARA_CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    stager_config = VectaraUploadStagerConfig(batch_size=10)
+    stager = VectaraUploadStager(upload_stager_config=stager_config)
+    new_upload_file = stager.run(
+        elements_filepath=upload_file,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+        file_data=file_data,
+    )
+
+    uploader = VectaraUploader(
+        connection_config=VectaraConnectionConfig(
+            **connection_kwargs,
+            access_config=VectaraAccessConfig(
+                oauth_client_id=oauth_client_id, oauth_secret=oauth_secret
+            ),
+        ),
+        upload_config=VectaraUploaderConfig(),
+    )
+
+    with new_upload_file.open() as new_upload_fp:
+        elements_stager = json.load(new_upload_fp)
+
+    if uploader.is_async():
+        await uploader.run_data_async(data=elements_stager, file_data=file_data)
+
+    with upload_file.open() as upload_fp:
+        elements = json.load(upload_fp)
+    first_element = elements[0]
+
+    for i in range(retries):
+        response = query_data(corpus_key, first_element["element_id"])
+        if not response["search_results"]:
+            time.sleep(interval)
+        else:
+            break
+
+    validate_upload(response=response, expected_data=first_element)

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[1] / "env_setup"
+expected_results_path = Path(__file__).parents[1] / "expected_results"

test/integration/connectors/utils/docker.py

@@ -0,0 +1,151 @@
+import time
+from contextlib import contextmanager
+from typing import Optional, Union
+
+import docker
+from docker.models.containers import Container
+from pydantic import BaseModel, Field, field_serializer
+
+
+class HealthCheck(BaseModel):
+    test: Union[str, list[str]]
+    interval: int = Field(
+        gt=0, default=30, description="The time to wait between checks in seconds."
+    )
+    timeout: int = Field(
+        gt=0, default=30, description="The time to wait before considering the check to have hung."
+    )
+    retries: int = Field(
+        gt=0,
+        default=3,
+        description="The number of consecutive failures needed "
+        "to consider a container as unhealthy.",
+    )
+    start_period: int = Field(
+        gt=0,
+        default=0,
+        description="Start period for the container to initialize before starting health-retries countdown in seconds.",  # noqa: E501
+    )
+
+    @field_serializer("interval")
+    def serialize_interval(self, interval: int) -> int:
+        return int(interval * 10e8)
+
+    @field_serializer("timeout")
+    def serialize_timeout(self, timeout: int) -> int:
+        return int(timeout * 10e8)
+
+    @field_serializer("start_period")
+    def serialize_start_period(self, start_period: int) -> int:
+        return int(start_period * 10e8)
+
+
+def get_container(
+    docker_client: docker.DockerClient,
+    image: str,
+    ports: dict,
+    name: Optional[str] = "connector_test",
+    environment: Optional[dict] = None,
+    volumes: Optional[dict] = None,
+    healthcheck: Optional[HealthCheck] = None,
+) -> Container:
+    run_kwargs = {
+        "image": image,
+        "detach": True,
+        "ports": ports,
+    }
+    if environment:
+        run_kwargs["environment"] = environment
+    if volumes:
+        run_kwargs["volumes"] = volumes
+    if healthcheck:
+        run_kwargs["healthcheck"] = healthcheck.model_dump()
+    if name:
+        run_kwargs["name"] = name
+    container: Container = docker_client.containers.run(**run_kwargs)
+    return container
+
+
+def get_healthcheck(container: Container) -> Optional[HealthCheck]:
+    healthcheck_config = container.attrs.get("Config", {}).get("Healthcheck", None)
+    if not healthcheck_config:
+        return None
+    healthcheck_data = {
+        "test": healthcheck_config["Test"],
+    }
+    if interval := healthcheck_config.get("Interval"):
+        healthcheck_data["interval"] = interval / 10e8
+    if start_period := healthcheck_config.get("StartPeriod"):
+        healthcheck_data["start_period"] = start_period / 10e8
+    if retries := healthcheck_config.get("Retries"):
+        healthcheck_data["retries"] = retries
+    return HealthCheck.model_validate(healthcheck_data)
+
+
+def healthcheck_wait(
+    container: Container, retries: int = 30, interval: int = 1, start_period: Optional[int] = None
+) -> None:
+    if start_period:
+        time.sleep(start_period)
+    health = container.health
+    tries = 0
+    while health != "healthy" and tries < retries:
+        tries += 1
+        logs = container.attrs.get("State", {}).get("Health", {}).get("Log")
+        latest_log = logs[-1] if logs else None
+        print(
+            f"attempt {tries} - waiting for docker container "
+            f"to be healthy: {health} latest log: {latest_log}"
+        )
+        time.sleep(interval)
+        container.reload()
+        health = container.health
+    if health != "healthy":
+        logs = container.attrs.get("State", {}).get("Health", {}).get("Log")
+        latest_log = logs[-1] if logs else None
+        raise TimeoutError(f"Docker container never came up healthy: {latest_log}")
+
+
+@contextmanager
+def container_context(
+    image: str,
+    ports: dict,
+    environment: Optional[dict] = None,
+    volumes: Optional[dict] = None,
+    healthcheck: Optional[HealthCheck] = None,
+    healthcheck_retries: int = 30,
+    docker_client: Optional[docker.DockerClient] = None,
+    name: Optional[str] = "connector_test",
+):
+    docker_client = docker_client or docker.from_env()
+    print(f"pulling image {image}")
+    docker_client.images.pull(image)
+    container: Optional[Container] = None
+    try:
+        container = get_container(
+            docker_client=docker_client,
+            image=image,
+            ports=ports,
+            environment=environment,
+            volumes=volumes,
+            healthcheck=healthcheck,
+            name=name,
+        )
+        if healthcheck_data := get_healthcheck(container):
+            # Mirror whatever healthcheck config set on container
+            healthcheck_wait(
+                container=container,
+                retries=healthcheck_retries,
+                start_period=healthcheck_data.start_period,
+                interval=healthcheck_data.interval,
+            )
+        yield container
+    except AssertionError as e:
+        if container:
+            logs = container.logs()
+            print(logs.decode("utf-8"))
+            raise e
+    finally:
+        if container:
+            container.kill()
+            container.remove()

test/integration/connectors/utils/docker_compose.py

@@ -0,0 +1,59 @@
+import subprocess
+from contextlib import contextmanager
+from pathlib import Path
+
+
+def docker_compose_down(docker_compose_path: Path):
+    cmd = f"docker compose -f {docker_compose_path.resolve()} down --remove-orphans -v --rmi all"
+    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")))
+
+
+def run_cleanup(docker_compose_path: Path):
+    docker_compose_down(docker_compose_path=docker_compose_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")))
+        cmd = f"docker compose -f {docker_compose_path.resolve()} logs"
+        logs = subprocess.run(
+            cmd,
+            shell=True,
+            capture_output=True,
+        )
+        print("DOCKER LOGS: {}".format(logs.stdout.decode("utf-8")))
+        raise e
+    finally:
+        run_cleanup(docker_compose_path=docker_compose_path)