unstructured-ingest 0.1.0__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/{test_postgres.py → sql/test_postgres.py}

@@ -1,21 +1,93 @@
 import tempfile
+from contextlib import contextmanager
 from pathlib import Path
 
 import pandas as pd
 import pytest
 from psycopg2 import connect
 
-from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
+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 import (
+    ValidationConfigs,
+    source_connector_validation,
+)
 from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.processes.connectors.sql.postgres import (
     CONNECTOR_TYPE,
     PostgresAccessConfig,
     PostgresConnectionConfig,
+    PostgresDownloader,
+    PostgresDownloaderConfig,
+    PostgresIndexer,
+    PostgresIndexerConfig,
     PostgresUploader,
     PostgresUploadStager,
 )
 
+SEED_DATA_ROWS = 20
+
+
+@contextmanager
+def postgres_download_setup() -> None:
+    with docker_compose_context(docker_compose_path=env_setup_path / "sql" / "postgres" / "source"):
+        connection = connect(
+            user="unstructured",
+            password="test",
+            dbname="test_db",
+            host="localhost",
+            port=5433,
+        )
+        with connection.cursor() as cursor:
+            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
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+async def test_postgres_source():
+    connect_params = {
+        "host": "localhost",
+        "port": 5433,
+        "database": "test_db",
+        "user": "unstructured",
+        "password": "test",
+    }
+    with postgres_download_setup():
+        with tempfile.TemporaryDirectory() as tmpdir:
+            connection_config = PostgresConnectionConfig(
+                host=connect_params["host"],
+                port=connect_params["port"],
+                database=connect_params["database"],
+                username=connect_params["user"],
+                access_config=PostgresAccessConfig(password=connect_params["password"]),
+            )
+            indexer = PostgresIndexer(
+                connection_config=connection_config,
+                index_config=PostgresIndexerConfig(
+                    table_name="cars", id_column="car_id", batch_size=5
+                ),
+            )
+            downloader = PostgresDownloader(
+                connection_config=connection_config,
+                download_config=PostgresDownloaderConfig(
+                    fields=["car_id", "brand"], download_dir=Path(tmpdir)
+                ),
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="postgres",
+                    expected_num_files=SEED_DATA_ROWS,
+                    expected_number_indexed_file_data=4,
+                    validate_downloaded_files=True,
+                ),
+            )
+
 
 def validate_destination(
     connect_params: dict,
@@ -50,7 +122,9 @@ async def test_postgres_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)
-    with docker_compose_context(docker_compose_path=env_setup_path / "sql"):
+    with docker_compose_context(
+        docker_compose_path=env_setup_path / "sql" / "postgres" / "destination"
+    ):
         with tempfile.TemporaryDirectory() as tmpdir:
             stager = PostgresUploadStager()
             stager_params = {

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/{test_sqlite.py → sql/test_sqlite.py}

@@ -6,36 +6,92 @@ from pathlib import Path
 import pandas as pd
 import pytest
 
-from test.integration.connectors.utils.constants import DESTINATION_TAG, env_setup_path
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SOURCE_TAG, env_setup_path
+from test.integration.connectors.utils.validation import (
+    ValidationConfigs,
+    source_connector_validation,
+)
 from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.processes.connectors.sql.sqlite import (
     CONNECTOR_TYPE,
     SQLiteConnectionConfig,
+    SQLiteDownloader,
+    SQLiteDownloaderConfig,
+    SQLiteIndexer,
+    SQLiteIndexerConfig,
     SQLiteUploader,
     SQLiteUploadStager,
 )
 
+SEED_DATA_ROWS = 20
+
+
+@contextmanager
+def sqlite_download_setup() -> Path:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = Path(tmpdir) / "mock_database.db"
+        db_init_path = env_setup_path / "sql" / "sqlite" / "source" / "sqlite-schema.sql"
+        assert db_init_path.exists()
+        assert db_init_path.is_file()
+        with sqlite3.connect(database=db_path) as sqlite_connection:
+            cursor = sqlite_connection.cursor()
+            with db_init_path.open("r") as f:
+                query = f.read()
+            cursor.executescript(query)
+            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()
+            cursor.close()
+        yield db_path
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+async def test_sqlite_source():
+    with sqlite_download_setup() as db_path:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            connection_config = SQLiteConnectionConfig(database_path=db_path)
+            indexer = SQLiteIndexer(
+                connection_config=connection_config,
+                index_config=SQLiteIndexerConfig(
+                    table_name="cars", id_column="car_id", batch_size=5
+                ),
+            )
+            downloader = SQLiteDownloader(
+                connection_config=connection_config,
+                download_config=SQLiteDownloaderConfig(
+                    fields=["car_id", "brand"], download_dir=Path(tmpdir)
+                ),
+            )
+            await source_connector_validation(
+                indexer=indexer,
+                downloader=downloader,
+                configs=ValidationConfigs(
+                    test_id="sqlite",
+                    expected_num_files=SEED_DATA_ROWS,
+                    expected_number_indexed_file_data=4,
+                    validate_downloaded_files=True,
+                ),
+            )
+
 
 @contextmanager
-def sqlite_setup() -> Path:
+def sqlite_upload_setup() -> Path:
     # Provision the local file that sqlite points to to have the desired schema for the integration
     # tests and make sure the file and connection get cleaned up by using a context manager.
     with tempfile.TemporaryDirectory() as tmpdir:
         db_path = Path(tmpdir) / "elements.db"
-        db_init_path = env_setup_path / "sql" / "sqlite-schema.sql"
+        db_init_path = env_setup_path / "sql" / "sqlite" / "destination" / "sqlite-schema.sql"
         assert db_init_path.exists()
         assert db_init_path.is_file()
-        connection = None
-        try:
-            connection = sqlite3.connect(database=db_path)
+        with sqlite3.connect(database=db_path) as sqlite_connection:
             with db_init_path.open("r") as f:
                 query = f.read()
-            cursor = connection.cursor()
+            cursor = sqlite_connection.cursor()
             cursor.executescript(query)
-            yield db_path
-        finally:
-            if connection:
-                connection.close()
+        yield db_path
 
 
 def validate_destination(db_path: Path, expected_num_elements: int):
@@ -62,7 +118,7 @@ async def test_sqlite_destination(upload_file: Path):
     # the sqlite 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)
-    with sqlite_setup() as db_path:
+    with sqlite_upload_setup() as db_path:
         with tempfile.TemporaryDirectory() as tmpdir:
             stager = SQLiteUploadStager()
             stager_params = {

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/constants.py

@@ -3,5 +3,5 @@ from pathlib import Path
 SOURCE_TAG = "source"
 DESTINATION_TAG = "destination"
 
-env_setup_path = Path(__file__).parents[4] / "test_e2e" / "env_setup"
+env_setup_path = Path(__file__).parents[1] / "env_setup"
 expected_results_path = Path(__file__).parents[1] / "expected_results"