unstructured-ingest 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

test/integration/connectors/test_kafka.py

@@ -1,11 +1,14 @@
-import socket
+import json
 import tempfile
+import time
 from pathlib import Path
 
 import pytest
-from confluent_kafka import Producer
+from confluent_kafka import Consumer, KafkaError, KafkaException, Producer
+from confluent_kafka.admin import AdminClient, NewTopic
 
 from test.integration.connectors.utils.constants import (
+    DESTINATION_TAG,
     SOURCE_TAG,
     env_setup_path,
 )
@@ -14,6 +17,8 @@ from test.integration.connectors.utils.validation import (
     ValidationConfigs,
     source_connector_validation,
 )
+from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.processes.connectors.kafka.local import (
     CONNECTOR_TYPE,
     LocalKafkaConnectionConfig,
@@ -21,27 +26,64 @@ from unstructured_ingest.v2.processes.connectors.kafka.local import (
     LocalKafkaDownloaderConfig,
     LocalKafkaIndexer,
     LocalKafkaIndexerConfig,
+    LocalKafkaUploader,
+    LocalKafkaUploaderConfig,
 )
 
 SEED_MESSAGES = 10
 TOPIC = "fake-topic"
 
 
+def get_admin_client() -> AdminClient:
+    conf = {
+        "bootstrap.servers": "localhost:29092",
+    }
+    return AdminClient(conf)
+
+
 @pytest.fixture
-def kafka_seed_topic() -> str:
-    with docker_compose_context(docker_compose_path=env_setup_path / "kafka"):
-        conf = {
-            "bootstrap.servers": "localhost:29092",
-            "client.id": socket.gethostname(),
-            "message.max.bytes": 10485760,
-        }
-        producer = Producer(conf)
-        for i in range(SEED_MESSAGES):
-            message = f"This is some text for message {i}"
-            producer.produce(topic=TOPIC, value=message)
-        producer.flush(timeout=10)
-        print(f"kafka topic {TOPIC} seeded with {SEED_MESSAGES} messages")
-        yield TOPIC
+def docker_compose_ctx():
+    with docker_compose_context(docker_compose_path=env_setup_path / "kafka") as ctx:
+        yield ctx
+
+
+def wait_for_topic(topic: str, retries: int = 10, interval: int = 1):
+    admin_client = get_admin_client()
+    current_topics = admin_client.list_topics().topics
+    attempts = 0
+    while topic not in current_topics and attempts < retries:
+        attempts += 1
+        print(
+            "Attempt {}: Waiting for topic {} to exist in {}".format(
+                attempts, topic, ", ".join(current_topics)
+            )
+        )
+        time.sleep(interval)
+        current_topics = admin_client.list_topics().topics
+    if topic not in current_topics:
+        raise TimeoutError(f"Timeout out waiting for topic {topic} to exist")
+
+
+@pytest.fixture
+def kafka_seed_topic(docker_compose_ctx) -> str:
+    conf = {
+        "bootstrap.servers": "localhost:29092",
+    }
+    producer = Producer(conf)
+    for i in range(SEED_MESSAGES):
+        message = f"This is some text for message {i}"
+        producer.produce(topic=TOPIC, value=message)
+    producer.flush(timeout=10)
+    print(f"kafka topic {TOPIC} seeded with {SEED_MESSAGES} messages")
+    wait_for_topic(topic=TOPIC)
+    return TOPIC
+
+
+@pytest.fixture
+def kafka_upload_topic(docker_compose_ctx) -> str:
+    admin_client = get_admin_client()
+    admin_client.create_topics([NewTopic(TOPIC, 1, 1)])
+    return TOPIC
 
 
 @pytest.mark.asyncio
@@ -58,6 +100,7 @@ async def test_kafka_source_local(kafka_seed_topic: str):
         downloader = LocalKafkaDownloader(
             connection_config=connection_config, download_config=download_config
         )
+        indexer.precheck()
         await source_connector_validation(
             indexer=indexer,
             downloader=downloader,
@@ -65,3 +108,94 @@ async def test_kafka_source_local(kafka_seed_topic: str):
                 test_id="kafka", expected_num_files=5, validate_downloaded_files=True
             ),
         )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+def test_kafka_source_local_precheck_fail_no_cluster():
+    connection_config = LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092)
+    indexer = LocalKafkaIndexer(
+        connection_config=connection_config,
+        index_config=LocalKafkaIndexerConfig(topic=TOPIC, num_messages_to_consume=5),
+    )
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+def test_kafka_source_local_precheck_fail_no_topic(kafka_seed_topic: str):
+    connection_config = LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092)
+    indexer = LocalKafkaIndexer(
+        connection_config=connection_config,
+        index_config=LocalKafkaIndexerConfig(topic="topic", num_messages_to_consume=5),
+    )
+    with pytest.raises(SourceConnectionError):
+        indexer.precheck()
+
+
+def get_all_messages(topic: str, max_empty_messages: int = 5) -> list[dict]:
+    conf = {
+        "bootstrap.servers": "localhost:29092",
+        "group.id": "default_group_id",
+        "enable.auto.commit": "false",
+        "auto.offset.reset": "earliest",
+    }
+    consumer = Consumer(conf)
+    consumer.subscribe([topic])
+    messages = []
+    try:
+        empty_count = 0
+        while empty_count < max_empty_messages:
+            msg = consumer.poll(timeout=1)
+            if msg is None:
+                empty_count += 1
+                continue
+            if msg.error():
+                if msg.error().code() == KafkaError._PARTITION_EOF:
+                    break
+                else:
+                    raise KafkaException(msg.error())
+            try:
+                message = json.loads(msg.value().decode("utf8"))
+                messages.append(message)
+            finally:
+                consumer.commit(asynchronous=False)
+    finally:
+        print("closing consumer")
+        consumer.close()
+    return messages
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+async def test_kafka_destination_local(upload_file: Path, kafka_upload_topic: str):
+    uploader = LocalKafkaUploader(
+        connection_config=LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092),
+        upload_config=LocalKafkaUploaderConfig(topic=TOPIC, batch_size=10),
+    )
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    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)
+    all_messages = get_all_messages(topic=kafka_upload_topic)
+    with upload_file.open("r") as upload_fs:
+        content_to_upload = json.load(upload_fs)
+    assert len(all_messages) == len(content_to_upload), (
+        f"expected number of messages ({len(content_to_upload)}) doesn't "
+        f"match how many messages read off of kakfa topic {kafka_upload_topic}: {len(all_messages)}"
+    )
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+def test_kafka_destination_local_precheck_fail_no_cluster():
+    uploader = LocalKafkaUploader(
+        connection_config=LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092),
+        upload_config=LocalKafkaUploaderConfig(topic=TOPIC, batch_size=10),
+    )
+    with pytest.raises(DestinationConnectionError):
+        uploader.precheck()

test/integration/connectors/test_lancedb.py

@@ -0,0 +1,209 @@
+import os
+from pathlib import Path
+from typing import Literal, Union
+
+import lancedb
+import pandas as pd
+import pyarrow as pa
+import pytest
+import pytest_asyncio
+from lancedb import AsyncConnection
+from upath import UPath
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.lancedb.aws import (
+    LanceDBS3AccessConfig,
+    LanceDBS3ConnectionConfig,
+    LanceDBS3Uploader,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.azure import (
+    LanceDBAzureAccessConfig,
+    LanceDBAzureConnectionConfig,
+    LanceDBAzureUploader,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.gcp import (
+    LanceDBGCSAccessConfig,
+    LanceDBGCSConnectionConfig,
+    LanceDBGSPUploader,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.lancedb import (
+    CONNECTOR_TYPE,
+    LanceDBUploaderConfig,
+    LanceDBUploadStager,
+)
+from unstructured_ingest.v2.processes.connectors.lancedb.local import (
+    LanceDBLocalAccessConfig,
+    LanceDBLocalConnectionConfig,
+    LanceDBLocalUploader,
+)
+
+DATABASE_NAME = "database"
+TABLE_NAME = "elements"
+DIMENSION = 384
+NUMBER_EXPECTED_ROWS = 22
+NUMBER_EXPECTED_COLUMNS = 10
+S3_BUCKET = "s3://utic-ingest-test-fixtures/"
+GS_BUCKET = "gs://utic-test-ingest-fixtures-output/"
+AZURE_BUCKET = "az://utic-ingest-test-fixtures-output/"
+REQUIRED_ENV_VARS = {
+    "s3": ("S3_INGEST_TEST_ACCESS_KEY", "S3_INGEST_TEST_SECRET_KEY"),
+    "gcs": ("GCP_INGEST_SERVICE_KEY",),
+    "az": ("AZURE_DEST_CONNECTION_STR",),
+    "local": (),
+}
+
+
+SCHEMA = pa.schema(
+    [
+        pa.field("vector", pa.list_(pa.float16(), DIMENSION)),
+        pa.field("text", pa.string(), nullable=True),
+        pa.field("type", pa.string(), nullable=True),
+        pa.field("element_id", pa.string(), nullable=True),
+        pa.field("metadata-text_as_html", pa.string(), nullable=True),
+        pa.field("metadata-filetype", pa.string(), nullable=True),
+        pa.field("metadata-filename", pa.string(), nullable=True),
+        pa.field("metadata-languages", pa.list_(pa.string()), nullable=True),
+        pa.field("metadata-is_continuation", pa.bool_(), nullable=True),
+        pa.field("metadata-page_number", pa.int32(), nullable=True),
+    ]
+)
+
+
+@pytest_asyncio.fixture
+async def connection_with_uri(request, tmp_path: Path):
+    target = request.param
+    uri = _get_uri(target, local_base_path=tmp_path)
+
+    unset_variables = [env for env in REQUIRED_ENV_VARS[target] if env not in os.environ]
+    if unset_variables:
+        pytest.skip(
+            reason="Following required environment variables were not set: "
+            + f"{', '.join(unset_variables)}"
+        )
+
+    storage_options = {
+        "aws_access_key_id": os.getenv("S3_INGEST_TEST_ACCESS_KEY"),
+        "aws_secret_access_key": os.getenv("S3_INGEST_TEST_SECRET_KEY"),
+        "google_service_account_key": os.getenv("GCP_INGEST_SERVICE_KEY"),
+    }
+    azure_connection_string = os.getenv("AZURE_DEST_CONNECTION_STR")
+    if azure_connection_string:
+        storage_options.update(_parse_azure_connection_string(azure_connection_string))
+
+    storage_options = {key: value for key, value in storage_options.items() if value is not None}
+    connection = await lancedb.connect_async(
+        uri=uri,
+        storage_options=storage_options,
+    )
+    await connection.create_table(name=TABLE_NAME, schema=SCHEMA)
+
+    yield connection, uri
+
+    await connection.drop_database()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+@pytest.mark.parametrize("connection_with_uri", ["local", "s3", "gcs", "az"], indirect=True)
+async def test_lancedb_destination(
+    upload_file: Path,
+    connection_with_uri: tuple[AsyncConnection, str],
+    tmp_path: Path,
+) -> None:
+    connection, uri = connection_with_uri
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    stager = LanceDBUploadStager()
+    uploader = _get_uploader(uri)
+    staged_file_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+
+    await uploader.run_async(path=staged_file_path, file_data=file_data)
+
+    table = await connection.open_table(TABLE_NAME)
+    table_df: pd.DataFrame = await table.to_pandas()
+
+    assert len(table_df) == NUMBER_EXPECTED_ROWS
+    assert len(table_df.columns) == NUMBER_EXPECTED_COLUMNS
+
+    assert table_df["element_id"][0] == "2470d8dc42215b3d68413b55bf00fed2"
+    assert table_df["type"][0] == "CompositeElement"
+    assert table_df["metadata-filename"][0] == "DA-1p-with-duplicate-pages.pdf.json"
+    assert table_df["metadata-text_as_html"][0] is None
+
+
+def _get_uri(target: Literal["local", "s3", "gcs", "az"], local_base_path: Path) -> str:
+    if target == "local":
+        return str(local_base_path / DATABASE_NAME)
+    if target == "s3":
+        base_uri = UPath(S3_BUCKET)
+    elif target == "gcs":
+        base_uri = UPath(GS_BUCKET)
+    elif target == "az":
+        base_uri = UPath(AZURE_BUCKET)
+
+    return str(base_uri / "destination" / "lancedb" / DATABASE_NAME)
+
+
+def _get_uploader(
+    uri: str,
+) -> Union[LanceDBAzureUploader, LanceDBAzureUploader, LanceDBS3Uploader, LanceDBGSPUploader]:
+    target = uri.split("://", maxsplit=1)[0] if uri.startswith(("s3", "az", "gs")) else "local"
+    if target == "az":
+        azure_connection_string = os.getenv("AZURE_DEST_CONNECTION_STR")
+        access_config_kwargs = _parse_azure_connection_string(azure_connection_string)
+        return LanceDBAzureUploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBAzureConnectionConfig(
+                access_config=LanceDBAzureAccessConfig(**access_config_kwargs),
+                uri=uri,
+            ),
+        )
+
+    elif target == "s3":
+        return LanceDBS3Uploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBS3ConnectionConfig(
+                access_config=LanceDBS3AccessConfig(
+                    aws_access_key_id=os.getenv("S3_INGEST_TEST_ACCESS_KEY"),
+                    aws_secret_access_key=os.getenv("S3_INGEST_TEST_SECRET_KEY"),
+                ),
+                uri=uri,
+            ),
+        )
+    elif target == "gs":
+        return LanceDBGSPUploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBGCSConnectionConfig(
+                access_config=LanceDBGCSAccessConfig(
+                    google_service_account_key=os.getenv("GCP_INGEST_SERVICE_KEY")
+                ),
+                uri=uri,
+            ),
+        )
+    else:
+        return LanceDBLocalUploader(
+            upload_config=LanceDBUploaderConfig(table_name=TABLE_NAME),
+            connection_config=LanceDBLocalConnectionConfig(
+                access_config=LanceDBLocalAccessConfig(),
+                uri=uri,
+            ),
+        )
+
+
+def _parse_azure_connection_string(
+    connection_str: str,
+) -> dict[Literal["azure_storage_account_name", "azure_storage_account_key"], str]:
+    parameters = dict(keyvalue.split("=", maxsplit=1) for keyvalue in connection_str.split(";"))
+    return {
+        "azure_storage_account_name": parameters.get("AccountName"),
+        "azure_storage_account_key": parameters.get("AccountKey"),
+    }

test/integration/connectors/test_milvus.py

@@ -0,0 +1,141 @@
+import json
+import time
+from pathlib import Path
+
+import docker
+import pytest
+from pymilvus import (
+    CollectionSchema,
+    DataType,
+    FieldSchema,
+    MilvusClient,
+)
+from pymilvus.milvus_client import IndexParams
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
+from test.integration.connectors.utils.docker import healthcheck_wait
+from test.integration.connectors.utils.docker_compose import docker_compose_context
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.milvus import (
+    CONNECTOR_TYPE,
+    MilvusConnectionConfig,
+    MilvusUploader,
+    MilvusUploaderConfig,
+    MilvusUploadStager,
+)
+
+DB_URI = "http://localhost:19530"
+DB_NAME = "test_database"
+COLLECTION_NAME = "test_collection"
+
+
+def get_schema() -> CollectionSchema:
+    id_field = FieldSchema(
+        name="id", dtype=DataType.INT64, description="primary field", is_primary=True, auto_id=True
+    )
+    embeddings_field = FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=384)
+    record_id_field = FieldSchema(name="record_id", dtype=DataType.VARCHAR, max_length=64)
+
+    schema = CollectionSchema(
+        enable_dynamic_field=True,
+        fields=[
+            id_field,
+            record_id_field,
+            embeddings_field,
+        ],
+    )
+
+    return schema
+
+
+def get_index_params() -> IndexParams:
+    index_params = IndexParams()
+    index_params.add_index(field_name="embeddings", index_type="AUTOINDEX", metric_type="COSINE")
+    index_params.add_index(field_name="record_id", index_type="Trie")
+    return index_params
+
+
+@pytest.fixture
+def collection():
+    docker_client = docker.from_env()
+    with docker_compose_context(docker_compose_path=env_setup_path / "milvus"):
+        milvus_container = docker_client.containers.get("milvus-standalone")
+        healthcheck_wait(container=milvus_container)
+        milvus_client = MilvusClient(uri=DB_URI)
+        try:
+            # Create the database
+            database_resp = milvus_client._get_connection().create_database(db_name=DB_NAME)
+            milvus_client.using_database(db_name=DB_NAME)
+
+            print(f"Created database {DB_NAME}: {database_resp}")
+
+            # Create the collection
+            schema = get_schema()
+            index_params = get_index_params()
+            collection_resp = milvus_client.create_collection(
+                collection_name=COLLECTION_NAME, schema=schema, index_params=index_params
+            )
+            print(f"Created collection {COLLECTION_NAME}: {collection_resp}")
+            yield COLLECTION_NAME
+        finally:
+            milvus_client.close()
+
+
+def get_count(client: MilvusClient) -> int:
+    count_field = "count(*)"
+    resp = client.query(collection_name="test_collection", output_fields=[count_field])
+    return resp[0][count_field]
+
+
+def validate_count(
+    client: MilvusClient, expected_count: int, retries: int = 10, interval: int = 1
+) -> None:
+    current_count = get_count(client=client)
+    retry_count = 0
+    while current_count != expected_count and retry_count < retries:
+        time.sleep(interval)
+        current_count = get_count(client=client)
+        retry_count += 1
+    assert current_count == expected_count, (
+        f"Expected count ({expected_count}) doesn't match how "
+        f"much came back from collection: {current_count}"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+async def test_milvus_destination(
+    upload_file: Path,
+    collection: str,
+    tmp_path: Path,
+):
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+    stager = MilvusUploadStager()
+    uploader = MilvusUploader(
+        connection_config=MilvusConnectionConfig(uri=DB_URI),
+        upload_config=MilvusUploaderConfig(collection_name=collection, db_name=DB_NAME),
+    )
+    staged_filepath = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+    uploader.precheck()
+    uploader.run(path=staged_filepath, file_data=file_data)
+
+    # Run validation
+    with staged_filepath.open() as f:
+        staged_elements = json.load(f)
+    expected_count = len(staged_elements)
+    with uploader.get_client() as client:
+        validate_count(client=client, expected_count=expected_count)
+
+    # Rerun and make sure the same documents get updated
+    uploader.run(path=staged_filepath, file_data=file_data)
+    with uploader.get_client() as client:
+        validate_count(client=client, expected_count=expected_count)