unstructured-ingest 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

test/integration/connectors/conftest.py

@@ -1,7 +1,11 @@
+import tempfile
 from pathlib import Path
+from typing import Generator
 
 import pytest
 
+from unstructured_ingest.v2.logger import logger
+
 FILENAME = "DA-1p-with-duplicate-pages.pdf.json"
 
 
@@ -13,3 +17,12 @@ def upload_file() -> Path:
     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_tests/test_volumes_native.py

@@ -16,7 +16,7 @@ from test.integration.connectors.utils.validation import (
     source_connector_validation,
 )
 from test.integration.utils import requires_env
-from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
 from unstructured_ingest.v2.processes.connectors.databricks.volumes_native import (
     CONNECTOR_TYPE,
     DatabricksNativeVolumesAccessConfig,
@@ -139,7 +139,11 @@ def validate_upload(client: WorkspaceClient, catalog: str, volume: str, volume_p
 async def test_volumes_native_destination(upload_file: Path):
     env_data = get_env_data()
     volume_path = f"databricks-volumes-test-output-{uuid.uuid4()}"
-    mock_file_data = FileData(identifier="mock file data", connector_type=CONNECTOR_TYPE)
+    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:
@@ -153,9 +157,9 @@ async def test_volumes_native_destination(upload_file: Path):
             ),
         )
         if uploader.is_async():
-            await uploader.run_async(path=upload_file, file_data=mock_file_data)
+            await uploader.run_async(path=upload_file, file_data=file_data)
         else:
-            uploader.run(path=upload_file, file_data=mock_file_data)
+            uploader.run(path=upload_file, file_data=file_data)
 
         validate_upload(
             client=workspace_client,

test/integration/connectors/sql/test_postgres.py

@@ -2,7 +2,6 @@ import tempfile
 from contextlib import contextmanager
 from pathlib import Path
 
-import faker
 import pandas as pd
 import pytest
 from psycopg2 import connect
@@ -26,9 +25,7 @@ from unstructured_ingest.v2.processes.connectors.sql.postgres import (
     PostgresUploadStager,
 )
 
-faker = faker.Faker()
-
-SEED_DATA_ROWS = 40
+SEED_DATA_ROWS = 20
 
 
 @contextmanager
@@ -42,11 +39,8 @@ def postgres_download_setup() -> None:
             port=5433,
         )
         with connection.cursor() as cursor:
-            for _ in range(SEED_DATA_ROWS):
-                sql_statment = (
-                    f"INSERT INTO cars (brand, price) VALUES "
-                    f"('{faker.word()}', {faker.random_int()})"
-                )
+            for i in range(SEED_DATA_ROWS):
+                sql_statment = f"INSERT INTO cars (brand, price) VALUES " f"('brand_{i}', {i})"
                 cursor.execute(sql_statment)
             connection.commit()
         yield
@@ -88,7 +82,9 @@ async def test_postgres_source():
                 downloader=downloader,
                 configs=ValidationConfigs(
                     test_id="postgres",
-                    expected_num_files=40,
+                    expected_num_files=SEED_DATA_ROWS,
+                    expected_number_indexed_file_data=4,
+                    validate_downloaded_files=True,
                 ),
             )
 

test/integration/connectors/sql/test_snowflake.py

@@ -0,0 +1,205 @@
+import os
+import tempfile
+from pathlib import Path
+
+import docker
+import pandas as pd
+import pytest
+import snowflake.connector as sf
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG, env_setup_path
+from test.integration.connectors.utils.docker import container_context
+from test.integration.connectors.utils.validation import (
+    ValidationConfigs,
+    source_connector_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.processes.connectors.sql.snowflake import (
+    CONNECTOR_TYPE,
+    SnowflakeAccessConfig,
+    SnowflakeConnectionConfig,
+    SnowflakeDownloader,
+    SnowflakeDownloaderConfig,
+    SnowflakeIndexer,
+    SnowflakeIndexerConfig,
+    SnowflakeUploader,
+    SnowflakeUploadStager,
+)
+
+SEED_DATA_ROWS = 20
+
+
+def seed_data():
+    conn = sf.connect(
+        user="test",
+        password="test",
+        account="test",
+        database="test",
+        host="snowflake.localhost.localstack.cloud",
+    )
+
+    file = Path(env_setup_path / "sql" / "snowflake" / "source" / "snowflake-schema.sql")
+
+    with file.open() as f:
+        sql = f.read()
+
+    cur = conn.cursor()
+    cur.execute(sql)
+    for i in range(SEED_DATA_ROWS):
+        sql_statment = f"INSERT INTO cars (brand, price) VALUES " f"('brand_{i}', {i})"
+        cur.execute(sql_statment)
+
+    cur.close()
+    conn.close()
+
+
+def init_db_destination():
+    conn = sf.connect(
+        user="test",
+        password="test",
+        account="test",
+        database="test",
+        host="snowflake.localhost.localstack.cloud",
+    )
+
+    file = Path(env_setup_path / "sql" / "snowflake" / "destination" / "snowflake-schema.sql")
+
+    with file.open() as f:
+        sql = f.read()
+
+    cur = conn.cursor()
+    cur.execute(sql)
+
+    cur.close()
+    conn.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+@requires_env("LOCALSTACK_AUTH_TOKEN")
+async def test_snowflake_source():
+    docker_client = docker.from_env()
+    token = os.getenv("LOCALSTACK_AUTH_TOKEN")
+    with container_context(
+        docker_client=docker_client,
+        image="localstack/snowflake",
+        environment={"LOCALSTACK_AUTH_TOKEN": token, "EXTRA_CORS_ALLOWED_ORIGINS": "*"},
+        ports={4566: 4566, 443: 443},
+        healthcheck_timeout=30,
+    ):
+        seed_data()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            connection_config = SnowflakeConnectionConfig(
+                access_config=SnowflakeAccessConfig(password="test"),
+                account="test",
+                user="test",
+                database="test",
+                host="snowflake.localhost.localstack.cloud",
+            )
+            indexer = SnowflakeIndexer(
+                connection_config=connection_config,
+                index_config=SnowflakeIndexerConfig(
+                    table_name="cars", id_column="CAR_ID", batch_size=5
+                ),
+            )
+            downloader = SnowflakeDownloader(
+                connection_config=connection_config,
+                download_config=SnowflakeDownloaderConfig(
+                    fields=["CAR_ID", "BRAND"], download_dir=Path(tmpdir)
+                ),
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="snowflake",
+                    expected_num_files=SEED_DATA_ROWS,
+                    expected_number_indexed_file_data=4,
+                    validate_downloaded_files=True,
+                ),
+            )
+
+
+def validate_destination(
+    connect_params: dict,
+    expected_num_elements: int,
+):
+    # Run the following validations:
+    # * Check that the number of records in the table match the expected value
+    # * Given the embedding, make sure it matches the associated text it belongs to
+    conn = sf.connect(**connect_params)
+    cursor = conn.cursor()
+    try:
+        query = "select count(*) from elements;"
+        cursor.execute(query)
+        count = cursor.fetchone()[0]
+        assert (
+            count == expected_num_elements
+        ), f"dest check failed: got {count}, expected {expected_num_elements}"
+    finally:
+        cursor.close()
+        conn.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "sql")
+@requires_env("LOCALSTACK_AUTH_TOKEN")
+async def test_snowflake_destination(upload_file: Path):
+    # the postgres destination connector doesn't leverage the file data but is required as an input,
+    # mocking it with arbitrary values to meet the base requirements:
+    mock_file_data = FileData(identifier="mock file data", connector_type=CONNECTOR_TYPE)
+    docker_client = docker.from_env()
+    token = os.getenv("LOCALSTACK_AUTH_TOKEN")
+    with container_context(
+        docker_client=docker_client,
+        image="localstack/snowflake",
+        environment={"LOCALSTACK_AUTH_TOKEN": token, "EXTRA_CORS_ALLOWED_ORIGINS": "*"},
+        ports={4566: 4566, 443: 443},
+        healthcheck_timeout=30,
+    ):
+        init_db_destination()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            stager = SnowflakeUploadStager()
+            stager_params = {
+                "elements_filepath": upload_file,
+                "file_data": mock_file_data,
+                "output_dir": Path(tmpdir),
+                "output_filename": "test_db",
+            }
+            if stager.is_async():
+                staged_path = await stager.run_async(**stager_params)
+            else:
+                staged_path = stager.run(**stager_params)
+
+            # The stager should append the `.json` suffix to the output filename passed in.
+            assert staged_path.name == "test_db.json"
+
+            connect_params = {
+                "user": "test",
+                "password": "test",
+                "account": "test",
+                "database": "test",
+                "host": "snowflake.localhost.localstack.cloud",
+            }
+
+            uploader = SnowflakeUploader(
+                connection_config=SnowflakeConnectionConfig(
+                    access_config=SnowflakeAccessConfig(password=connect_params["password"]),
+                    account=connect_params["account"],
+                    user=connect_params["user"],
+                    database=connect_params["database"],
+                    host=connect_params["host"],
+                )
+            )
+            if uploader.is_async():
+                await uploader.run_async(path=staged_path, file_data=mock_file_data)
+            else:
+                uploader.run(path=staged_path, file_data=mock_file_data)
+
+            staged_df = pd.read_json(staged_path, orient="records", lines=True)
+            expected_num_elements = len(staged_df)
+            validate_destination(
+                connect_params=connect_params,
+                expected_num_elements=expected_num_elements,
+            )

test/integration/connectors/sql/test_sqlite.py

@@ -3,7 +3,6 @@ import tempfile
 from contextlib import contextmanager
 from pathlib import Path
 
-import faker
 import pandas as pd
 import pytest
 
@@ -24,9 +23,7 @@ from unstructured_ingest.v2.processes.connectors.sql.sqlite import (
     SQLiteUploadStager,
 )
 
-faker = faker.Faker()
-
-SEED_DATA_ROWS = 40
+SEED_DATA_ROWS = 20
 
 
 @contextmanager
@@ -41,11 +38,8 @@ def sqlite_download_setup() -> Path:
             with db_init_path.open("r") as f:
                 query = f.read()
             cursor.executescript(query)
-            for _ in range(SEED_DATA_ROWS):
-                sql_statment = (
-                    f"INSERT INTO cars (brand, price) "
-                    f"VALUES ('{faker.word()}', {faker.random_int()})"
-                )
+            for i in range(SEED_DATA_ROWS):
+                sql_statment = f"INSERT INTO cars (brand, price) " f"VALUES ('brand{i}', {i})"
                 cursor.execute(sql_statment)
 
             sqlite_connection.commit()
@@ -76,7 +70,9 @@ async def test_sqlite_source():
                 downloader=downloader,
                 configs=ValidationConfigs(
                     test_id="sqlite",
-                    expected_num_files=40,
+                    expected_num_files=SEED_DATA_ROWS,
+                    expected_number_indexed_file_data=4,
+                    validate_downloaded_files=True,
                 ),
             )
 

test/integration/connectors/test_delta_table.py

@@ -0,0 +1,138 @@
+import multiprocessing
+import os
+from pathlib import Path
+
+import pytest
+from deltalake import DeltaTable
+from fsspec import get_filesystem_class
+
+from test.integration.connectors.utils.constants import (
+    DESTINATION_TAG,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.delta_table import (
+    CONNECTOR_TYPE,
+    DeltaTableAccessConfig,
+    DeltaTableConnectionConfig,
+    DeltaTableUploader,
+    DeltaTableUploaderConfig,
+    DeltaTableUploadStager,
+    DeltaTableUploadStagerConfig,
+)
+
+multiprocessing.set_start_method("spawn")
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG)
+async def test_delta_table_destination_local(upload_file: Path, temp_dir: Path):
+    destination_path = str(temp_dir)
+    connection_config = DeltaTableConnectionConfig(
+        access_config=DeltaTableAccessConfig(),
+        table_uri=destination_path,
+    )
+    stager_config = DeltaTableUploadStagerConfig()
+    stager = DeltaTableUploadStager(upload_stager_config=stager_config)
+    new_upload_file = stager.run(
+        elements_filepath=upload_file,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    upload_config = DeltaTableUploaderConfig()
+    uploader = DeltaTableUploader(connection_config=connection_config, upload_config=upload_config)
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(
+            fullpath=upload_file.name, filename=new_upload_file.name
+        ),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+
+    if uploader.is_async():
+        await uploader.run_async(path=new_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=new_upload_file, file_data=file_data)
+    delta_table_path = os.path.join(destination_path, upload_file.name)
+    delta_table = DeltaTable(table_uri=delta_table_path)
+    df = delta_table.to_pandas()
+
+    EXPECTED_COLUMNS = 10
+    EXPECTED_ROWS = 22
+    assert (
+        len(df) == EXPECTED_ROWS
+    ), f"Number of rows in table vs expected: {len(df)}/{EXPECTED_ROWS}"
+    assert (
+        len(df.columns) == EXPECTED_COLUMNS
+    ), f"Number of columns in table vs expected: {len(df.columns)}/{EXPECTED_COLUMNS}"
+
+
+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,
+        "AWS_REGION": "us-east-2",
+    }
+
+
+@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_delta_table_destination_s3(upload_file: Path, temp_dir: Path):
+    aws_credentials = get_aws_credentials()
+    s3_bucket = "s3://utic-platform-test-destination"
+    destination_path = f"{s3_bucket}/destination/test"
+    connection_config = DeltaTableConnectionConfig(
+        access_config=DeltaTableAccessConfig(
+            aws_access_key_id=aws_credentials["AWS_ACCESS_KEY_ID"],
+            aws_secret_access_key=aws_credentials["AWS_SECRET_ACCESS_KEY"],
+        ),
+        aws_region=aws_credentials["AWS_REGION"],
+        table_uri=destination_path,
+    )
+    stager_config = DeltaTableUploadStagerConfig()
+    stager = DeltaTableUploadStager(upload_stager_config=stager_config)
+    new_upload_file = stager.run(
+        elements_filepath=upload_file,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    upload_config = DeltaTableUploaderConfig()
+    uploader = DeltaTableUploader(connection_config=connection_config, upload_config=upload_config)
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(
+            fullpath=upload_file.name, filename=new_upload_file.name
+        ),
+        connector_type=CONNECTOR_TYPE,
+        identifier="mock file data",
+    )
+
+    try:
+        if uploader.is_async():
+            await uploader.run_async(path=new_upload_file, file_data=file_data)
+        else:
+            uploader.run(path=new_upload_file, file_data=file_data)
+        delta_table_path = os.path.join(destination_path, upload_file.name)
+        delta_table = DeltaTable(table_uri=delta_table_path, storage_options=aws_credentials)
+        df = delta_table.to_pandas()
+
+        EXPECTED_COLUMNS = 10
+        EXPECTED_ROWS = 22
+        assert (
+            len(df) == EXPECTED_ROWS
+        ), f"Number of rows in table vs expected: {len(df)}/{EXPECTED_ROWS}"
+        assert (
+            len(df.columns) == EXPECTED_COLUMNS
+        ), f"Number of columns in table vs expected: {len(df.columns)}/{EXPECTED_COLUMNS}"
+    finally:
+        s3fs = get_filesystem_class("s3")(
+            key=aws_credentials["AWS_ACCESS_KEY_ID"],
+            secret=aws_credentials["AWS_SECRET_ACCESS_KEY"],
+        )
+        s3fs.rm(path=destination_path, recursive=True)

test/integration/connectors/utils/docker.py

@@ -0,0 +1,78 @@
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+import docker
+from docker.models.containers import Container
+
+
+def get_container(
+    docker_client: docker.DockerClient,
+    image: str,
+    ports: dict,
+    environment: Optional[dict] = None,
+    volumes: Optional[dict] = None,
+    healthcheck: Optional[dict] = 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
+    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 healthcheck_wait(container: Container, timeout: int = 10) -> None:
+    health = container.health
+    start = time.time()
+    while health != "healthy" and time.time() - start < timeout:
+        time.sleep(1)
+        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}")
+
+
+@contextmanager
+def container_context(
+    docker_client: docker.DockerClient,
+    image: str,
+    ports: dict,
+    environment: Optional[dict] = None,
+    volumes: Optional[dict] = None,
+    healthcheck: Optional[dict] = None,
+    healthcheck_timeout: int = 10,
+):
+    container: Optional[Container] = None
+    try:
+        container = get_container(
+            docker_client=docker_client,
+            image=image,
+            ports=ports,
+            environment=environment,
+            volumes=volumes,
+            healthcheck=healthcheck,
+        )
+        if has_healthcheck(container):
+            healthcheck_wait(container=container, timeout=healthcheck_timeout)
+        yield container
+    except AssertionError as e:
+        if container:
+            logs = container.logs()
+            print(logs.decode("utf-8"))
+            raise e
+    finally:
+        if container:
+            container.kill()