unstructured-ingest 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

test/integration/connectors/test_mongodb.py

@@ -0,0 +1,332 @@
+import json
+import os
+import time
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator
+
+import pytest
+from pydantic import BaseModel, SecretStr
+from pymongo.collection import Collection
+from pymongo.database import Database
+from pymongo.mongo_client import MongoClient
+from pymongo.operations import SearchIndexModel
+
+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.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.mongodb import (
+    CONNECTOR_TYPE,
+    MongoDBAccessConfig,
+    MongoDBConnectionConfig,
+    MongoDBDownloader,
+    MongoDBDownloaderConfig,
+    MongoDBIndexer,
+    MongoDBIndexerConfig,
+    MongoDBUploader,
+    MongoDBUploaderConfig,
+)
+
+SOURCE_COLLECTION = "sample-mongodb-data"
+
+
+class EnvData(BaseModel):
+    uri: SecretStr
+    database: str
+
+
+def get_env_data() -> EnvData:
+    uri = os.getenv("MONGODB_URI")
+    assert uri
+    database = os.getenv("MONGODB_DATABASE")
+    assert database
+    return EnvData(uri=uri, database=database)
+
+
+@contextmanager
+def get_client() -> Generator[MongoClient, None, None]:
+    uri = get_env_data().uri.get_secret_value()
+    with MongoClient(uri) as client:
+        assert client.admin.command("ping")
+        yield client
+
+
+def wait_for_collection(
+    database: Database, collection_name: str, retries: int = 10, interval: int = 1
+):
+    collections = database.list_collection_names()
+    attempts = 0
+    while collection_name not in collections and attempts < retries:
+        attempts += 1
+        print(
+            "Waiting for collection {} to be recognized: {}".format(
+                collection_name, ", ".join(collections)
+            )
+        )
+        time.sleep(interval)
+        collections = database.list_collection_names()
+    if collection_name not in collection_name:
+        raise TimeoutError(f"Collection {collection_name} was not recognized")
+
+
+def get_search_index_status(collection: Collection, index_name: str) -> str:
+    search_indexes = collection.list_search_indexes(name=index_name)
+    search_index = list(search_indexes)[0]
+    return search_index["status"]
+
+
+def wait_for_search_index(
+    collection: Collection, index_name: str, retries: int = 60, interval: int = 1
+):
+    current_status = get_search_index_status(collection, index_name)
+    attempts = 0
+    while current_status != "READY" and attempts < retries:
+        attempts += 1
+        print(f"attempt {attempts}: waiting for search index to be READY: {current_status}")
+        time.sleep(interval)
+        current_status = get_search_index_status(collection, index_name)
+
+    if current_status != "READY":
+        raise TimeoutError("search index never detected as READY")
+
+
+@pytest.fixture
+def destination_collection() -> Collection:
+    env_data = get_env_data()
+    collection_name = f"utic-test-output-{uuid.uuid4()}"
+    with get_client() as client:
+        database = client[env_data.database]
+        print(f"creating collection in database {database}: {collection_name}")
+        collection = database.create_collection(name=collection_name)
+        search_index_name = "embeddings"
+        collection.create_search_index(
+            model=SearchIndexModel(
+                name=search_index_name,
+                definition={
+                    "mappings": {
+                        "dynamic": True,
+                        "fields": {
+                            "embeddings": [
+                                {"type": "knnVector", "dimensions": 384, "similarity": "euclidean"}
+                            ]
+                        },
+                    }
+                },
+            )
+        )
+        collection.create_index("record_id")
+        wait_for_collection(database=database, collection_name=collection_name)
+        wait_for_search_index(collection=collection, index_name=search_index_name)
+        try:
+            yield collection
+        finally:
+            print(f"deleting collection: {collection_name}")
+            collection.drop()
+
+
+def validate_collection_count(
+    collection: Collection, expected_records: int, retries: int = 10, interval: int = 1
+) -> None:
+    count = collection.count_documents(filter={})
+    attempt = 0
+    while count != expected_records and attempt < retries:
+        attempt += 1
+        print(f"attempt {attempt} to get count of collection {count} to match {expected_records}")
+        time.sleep(interval)
+        count = collection.count_documents(filter={})
+    assert (
+        count == expected_records
+    ), f"expected count ({expected_records}) does not match how many records were found: {count}"
+
+
+def validate_collection_vector(
+    collection: Collection, embedding: list[float], text: str, retries: int = 30, interval: int = 1
+) -> None:
+    pipeline = [
+        {
+            "$vectorSearch": {
+                "index": "embeddings",
+                "path": "embeddings",
+                "queryVector": embedding,
+                "numCandidates": 150,
+                "limit": 10,
+            },
+        },
+        {"$project": {"_id": 0, "text": 1, "score": {"$meta": "vectorSearchScore"}}},
+    ]
+    attempts = 0
+    results = list(collection.aggregate(pipeline=pipeline))
+    while not results and attempts < retries:
+        attempts += 1
+        print(f"attempt {attempts}, waiting for valid results: {results}")
+        time.sleep(interval)
+        results = list(collection.aggregate(pipeline=pipeline))
+    if not results:
+        raise TimeoutError("Timed out waiting for valid results")
+    print(f"found results on attempt {attempts}")
+    top_result = results[0]
+    assert top_result["score"] == 1.0, "score detected should be 1: {}".format(top_result["score"])
+    assert top_result["text"] == text, "text detected should be {}, found: {}".format(
+        text, top_result["text"]
+    )
+    for r in results[1:]:
+        assert r["score"] < 1.0, "score detected should be less than 1: {}".format(r["score"])
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("MONGODB_URI", "MONGODB_DATABASE")
+async def test_mongodb_source(temp_dir: Path):
+    env_data = get_env_data()
+    indexer_config = MongoDBIndexerConfig(database=env_data.database, collection=SOURCE_COLLECTION)
+    download_config = MongoDBDownloaderConfig(download_dir=temp_dir)
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
+    )
+    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
+    downloader = MongoDBDownloader(
+        connection_config=connection_config, download_config=download_config
+    )
+    await source_connector_validation(
+        indexer=indexer,
+        downloader=downloader,
+        configs=ValidationConfigs(
+            test_id=CONNECTOR_TYPE, expected_num_files=4, validate_downloaded_files=True
+        ),
+    )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+def test_mongodb_indexer_precheck_fail_no_host():
+    indexer_config = MongoDBIndexerConfig(
+        database="non-existent-database", collection="non-existent-database"
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri="mongodb+srv://ingest-test.hgaig.mongodb"),
+    )
+    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("MONGODB_URI", "MONGODB_DATABASE")
+def test_mongodb_indexer_precheck_fail_no_database():
+    env_data = get_env_data()
+    indexer_config = MongoDBIndexerConfig(
+        database="non-existent-database", collection=SOURCE_COLLECTION
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
+    )
+    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("MONGODB_URI", "MONGODB_DATABASE")
+def test_mongodb_indexer_precheck_fail_no_collection():
+    env_data = get_env_data()
+    indexer_config = MongoDBIndexerConfig(
+        database=env_data.database, collection="non-existent-collection"
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
+    )
+    indexer = MongoDBIndexer(connection_config=connection_config, index_config=indexer_config)
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("MONGODB_URI", "MONGODB_DATABASE")
+async def test_mongodb_destination(
+    upload_file: Path,
+    destination_collection: Collection,
+    tmp_path: Path,
+):
+    env_data = get_env_data()
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mongodb_mock_id",
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
+    )
+
+    upload_config = MongoDBUploaderConfig(
+        database=env_data.database,
+        collection=destination_collection.name,
+    )
+    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
+    uploader.precheck()
+    uploader.run(path=upload_file, file_data=file_data)
+
+    with upload_file.open() as f:
+        staged_elements = json.load(f)
+    expected_records = len(staged_elements)
+    validate_collection_count(collection=destination_collection, expected_records=expected_records)
+    first_element = staged_elements[0]
+    validate_collection_vector(
+        collection=destination_collection,
+        embedding=first_element["embeddings"],
+        text=first_element["text"],
+    )
+
+    uploader.run(path=upload_file, file_data=file_data)
+    validate_collection_count(collection=destination_collection, expected_records=expected_records)
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_mongodb_uploader_precheck_fail_no_host():
+    upload_config = MongoDBUploaderConfig(
+        database="database",
+        collection="collection",
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri="mongodb+srv://ingest-test.hgaig.mongodb"),
+    )
+    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
+    with pytest.raises(DestinationConnectionError):
+        uploader.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("MONGODB_URI", "MONGODB_DATABASE")
+def test_mongodb_uploader_precheck_fail_no_database():
+    env_data = get_env_data()
+    upload_config = MongoDBUploaderConfig(
+        database="database",
+        collection="collection",
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
+    )
+    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
+    with pytest.raises(DestinationConnectionError):
+        uploader.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@requires_env("MONGODB_URI", "MONGODB_DATABASE")
+def test_mongodb_uploader_precheck_fail_no_collection():
+    env_data = get_env_data()
+    upload_config = MongoDBUploaderConfig(
+        database=env_data.database,
+        collection="collection",
+    )
+    connection_config = MongoDBConnectionConfig(
+        access_config=MongoDBAccessConfig(uri=env_data.uri.get_secret_value()),
+    )
+    uploader = MongoDBUploader(connection_config=connection_config, upload_config=upload_config)
+    with pytest.raises(DestinationConnectionError):
+        uploader.precheck()

test/integration/connectors/test_pinecone.py

@@ -1,7 +1,9 @@
 import json
 import os
+import re
 import time
 from pathlib import Path
+from typing import Generator
 from uuid import uuid4
 
 import pytest
@@ -12,6 +14,7 @@ from test.integration.connectors.utils.constants import (
     DESTINATION_TAG,
 )
 from test.integration.utils import requires_env
+from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connectors.pinecone import (
@@ -24,6 +27,12 @@ from unstructured_ingest.v2.processes.connectors.pinecone import (
     PineconeUploadStagerConfig,
 )
 
+METADATA_BYTES_LIMIT = (
+    40960  # 40KB https://docs.pinecone.io/reference/quotas-and-limits#hard-limits
+)
+VECTOR_DIMENSION = 384
+SPEC = {"serverless": {"cloud": "aws", "region": "us-east-1"}}
+ALLOWED_METADATA_FIELD = "text"
 API_KEY = "PINECONE_API_KEY"
 
 
@@ -62,7 +71,7 @@ def wait_for_ready(client: Pinecone, index_name: str, timeout=60, interval=1) ->
 
 
 @pytest.fixture
-def pinecone_index() -> str:
+def pinecone_index() -> Generator[str, None, None]:
     pinecone = Pinecone(api_key=get_api_key())
     random_id = str(uuid4()).split("-")[0]
     index_name = f"ingest-test-{random_id}"
@@ -159,3 +168,46 @@ async def test_pinecone_destination(pinecone_index: str, upload_file: Path, temp
     validate_pinecone_index(
         index_name=pinecone_index, expected_num_of_vectors=expected_num_of_vectors
     )
+
+
+@requires_env(API_KEY)
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_large_metadata(pinecone_index: str, tmp_path: Path, upload_file: Path):
+    stager = PineconeUploadStager()
+    uploader = PineconeUploader(
+        connection_config=PineconeConnectionConfig(
+            access_config=PineconeAccessConfig(api_key=get_api_key()),
+            index_name=pinecone_index,
+        ),
+        upload_config=PineconeUploaderConfig(),
+    )
+    large_metadata_upload_file = tmp_path / "mock-upload-file.pdf.json"
+    large_metadata = {ALLOWED_METADATA_FIELD: "0" * 2 * METADATA_BYTES_LIMIT}
+
+    with open(upload_file) as file:
+        elements = json.load(file)
+
+    with open(large_metadata_upload_file, "w") as file:
+        mock_element = elements[0]
+        mock_element["metadata"] = large_metadata
+        json.dump([mock_element], file)
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(
+            fullpath=large_metadata_upload_file.name, filename=large_metadata_upload_file.name
+        ),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+    staged_file = stager.run(
+        file_data, large_metadata_upload_file, tmp_path, large_metadata_upload_file.name
+    )
+    try:
+        uploader.run(staged_file, file_data)
+    except DestinationConnectionError as e:
+        error_line = r"Metadata size is \d+ bytes, which exceeds the limit of \d+ bytes per vector"
+        if re.search(re.compile(error_line), str(e)) is None:
+            raise e
+        raise pytest.fail("Upload request failed due to metadata exceeding limits.")
+
+    validate_pinecone_index(pinecone_index, 1, interval=5)

test/integration/connectors/utils/docker.py

@@ -1,9 +1,43 @@
 import time
 from contextlib import contextmanager
-from typing import Optional
+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(
@@ -12,7 +46,7 @@ def get_container(
     ports: dict,
     environment: Optional[dict] = None,
     volumes: Optional[dict] = None,
-    healthcheck: Optional[dict] = None,
+    healthcheck: Optional[HealthCheck] = None,
 ) -> Container:
     run_kwargs = {
         "image": image,
@@ -24,25 +58,49 @@ def get_container(
     if volumes:
         run_kwargs["volumes"] = volumes
     if healthcheck:
-        run_kwargs["healthcheck"] = healthcheck
+        run_kwargs["healthcheck"] = healthcheck.model_dump()
     container: Container = docker_client.containers.run(**run_kwargs)
     return container
 
 
-def has_healthcheck(container: Container) -> bool:
-    return container.attrs.get("Config", {}).get("Healthcheck", None) is not None
+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, timeout: int = 10) -> None:
+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
-    start = time.time()
-    while health != "healthy" and time.time() - start < timeout:
-        time.sleep(1)
+    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":
-        health_dict = container.attrs.get("State", {}).get("Health", {})
-        raise TimeoutError(f"Docker container never came up healthy: {health_dict}")
+        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
@@ -51,11 +109,13 @@ def container_context(
     ports: dict,
     environment: Optional[dict] = None,
     volumes: Optional[dict] = None,
-    healthcheck: Optional[dict] = None,
-    healthcheck_timeout: int = 10,
+    healthcheck: Optional[HealthCheck] = None,
+    healthcheck_retries: int = 30,
     docker_client: Optional[docker.DockerClient] = None,
 ):
     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(
@@ -66,8 +126,14 @@ def container_context(
             volumes=volumes,
             healthcheck=healthcheck,
         )
-        if has_healthcheck(container):
-            healthcheck_wait(container=container, timeout=healthcheck_timeout)
+        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:

test/integration/connectors/utils/validation.py

@@ -240,6 +240,10 @@ def update_fixtures(
     # Rewrite the current file data
     if save_filedata:
         file_data_output_path = output_dir / "file_data"
+        print(
+            f"Writing {len(all_file_data)} file data to "
+            f"saved fixture location {file_data_output_path}"
+        )
         file_data_output_path.mkdir(parents=True, exist_ok=True)
         for file_data in all_file_data:
             file_data_path = file_data_output_path / f"{file_data.identifier}.json"
@@ -256,6 +260,10 @@ def update_fixtures(
     # If applicable, save raw downloads
     if save_downloads:
         raw_download_output_path = output_dir / "downloads"
+        print(
+            f"Writing {len(download_files)} downloaded files to "
+            f"saved fixture location {raw_download_output_path}"
+        )
         shutil.copytree(download_dir, raw_download_output_path)
 
 
@@ -328,6 +336,7 @@ async def source_connector_validation(
             postdownload_file_data = replace(resp["file_data"])
             all_postdownload_file_data.append(postdownload_file_data)
     if not overwrite_fixtures:
+        print("Running validation")
         run_all_validations(
             configs=configs,
             predownload_file_data=all_predownload_file_data,
@@ -336,6 +345,7 @@ async def source_connector_validation(
             test_output_dir=test_output_dir,
         )
     else:
+        print("Running fixtures update")
         update_fixtures(
             output_dir=test_output_dir,
             download_dir=download_dir,

test/integration/connectors/weaviate/conftest.py

@@ -0,0 +1,15 @@
+import json
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture
+def collections_schema_config() -> dict:
+    int_test_dir = Path(__file__).parent
+    assets_dir = int_test_dir / "assets"
+    config_file = assets_dir / "elements.json"
+    assert config_file.exists()
+    assert config_file.is_file()
+    with config_file.open() as config_data:
+        return json.load(config_data)