unstructured-ingest 0.3.13__py3-none-any.whl

test/integration/connectors/sql/test_postgres.py

@@ -0,0 +1,195 @@
+import json
+from pathlib import Path
+
+import pytest
+from _pytest.fixtures import TopRequest
+from psycopg2 import connect
+
+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.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
+    source_connector_validation,
+)
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.sql.postgres import (
+    CONNECTOR_TYPE,
+    PostgresAccessConfig,
+    PostgresConnectionConfig,
+    PostgresDownloader,
+    PostgresDownloaderConfig,
+    PostgresIndexer,
+    PostgresIndexerConfig,
+    PostgresUploader,
+    PostgresUploadStager,
+)
+
+SEED_DATA_ROWS = 10
+
+
+@pytest.fixture
+def source_database_setup() -> str:
+    db_name = "test_db"
+    with docker_compose_context(docker_compose_path=env_setup_path / "sql" / "postgres" / "source"):
+        connection = connect(
+            user="unstructured",
+            password="test",
+            dbname=db_name,
+            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 db_name
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+async def test_postgres_source(temp_dir: Path, source_database_setup: str):
+    connect_params = {
+        "host": "localhost",
+        "port": 5433,
+        "database": "test_db",
+        "user": "unstructured",
+        "password": "test",
+    }
+    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=6),
+    )
+    downloader = PostgresDownloader(
+        connection_config=connection_config,
+        download_config=PostgresDownloaderConfig(fields=["car_id", "brand"], download_dir=temp_dir),
+    )
+    await source_connector_validation(
+        indexer=indexer,
+        downloader=downloader,
+        configs=SourceValidationConfigs(
+            test_id="postgres",
+            expected_num_files=SEED_DATA_ROWS,
+            expected_number_indexed_file_data=2,
+            validate_downloaded_files=True,
+        ),
+    )
+
+
+def validate_destination(
+    connect_params: dict,
+    expected_num_elements: int,
+    test_embedding: list[float],
+    expected_text: str,
+):
+    # 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
+    with connect(**connect_params) as connection:
+        cursor = connection.cursor()
+        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}"
+
+        cursor.execute("SELECT embeddings FROM elements order by text limit 1")
+        similarity_query = (
+            f"SELECT text FROM elements ORDER BY embeddings <-> '{test_embedding}' LIMIT 1;"
+        )
+        cursor.execute(similarity_query)
+        res = cursor.fetchone()
+        assert res[0] == expected_text
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "sql")
+async def test_postgres_destination(upload_file: Path, temp_dir: 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,
+        source_identifiers=SourceIdentifiers(filename=upload_file.name, fullpath=upload_file.name),
+    )
+    with docker_compose_context(
+        docker_compose_path=env_setup_path / "sql" / "postgres" / "destination"
+    ):
+        stager = PostgresUploadStager()
+        staged_path = stager.run(
+            elements_filepath=upload_file,
+            file_data=mock_file_data,
+            output_dir=temp_dir,
+            output_filename=upload_file.name,
+        )
+
+        # The stager should append the `.json` suffix to the output filename passed in.
+        assert staged_path.suffix == upload_file.suffix
+
+        connect_params = {
+            "host": "localhost",
+            "port": 5433,
+            "database": "elements",
+            "user": "unstructured",
+            "password": "test",
+        }
+
+        uploader = PostgresUploader(
+            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"]),
+            )
+        )
+        uploader.precheck()
+        uploader.run(path=staged_path, file_data=mock_file_data)
+
+        with staged_path.open("r") as f:
+            staged_data = json.load(f)
+
+        sample_element = staged_data[0]
+        expected_num_elements = len(staged_data)
+        validate_destination(
+            connect_params=connect_params,
+            expected_num_elements=expected_num_elements,
+            expected_text=sample_element["text"],
+            test_embedding=sample_element["embeddings"],
+        )
+
+        uploader.run(path=staged_path, file_data=mock_file_data)
+        validate_destination(
+            connect_params=connect_params,
+            expected_num_elements=expected_num_elements,
+            expected_text=sample_element["text"],
+            test_embedding=sample_element["embeddings"],
+        )
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_postgres_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = PostgresUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/sql/test_singlestore.py

@@ -0,0 +1,176 @@
+import json
+from pathlib import Path
+
+import pytest
+import singlestoredb as s2
+from _pytest.fixtures import TopRequest
+
+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.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
+    source_connector_validation,
+)
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.sql.singlestore import (
+    CONNECTOR_TYPE,
+    SingleStoreAccessConfig,
+    SingleStoreConnectionConfig,
+    SingleStoreDownloader,
+    SingleStoreDownloaderConfig,
+    SingleStoreIndexer,
+    SingleStoreIndexerConfig,
+    SingleStoreUploader,
+    SingleStoreUploaderConfig,
+    SingleStoreUploadStager,
+)
+
+SEED_DATA_ROWS = 10
+
+
+@pytest.fixture
+def source_database_setup() -> dict:
+    connect_params = {
+        "host": "localhost",
+        "port": 3306,
+        "database": "ingest_test",
+        "user": "root",
+        "password": "password",
+    }
+    with docker_compose_context(
+        docker_compose_path=env_setup_path / "sql" / "singlestore" / "source"
+    ):
+        with s2.connect(**connect_params) as connection:
+            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 connect_params
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+async def test_singlestore_source(temp_dir: Path, source_database_setup: dict):
+
+    connection_config = SingleStoreConnectionConfig(
+        host=source_database_setup["host"],
+        port=source_database_setup["port"],
+        database=source_database_setup["database"],
+        user=source_database_setup["user"],
+        access_config=SingleStoreAccessConfig(password=source_database_setup["password"]),
+    )
+    indexer = SingleStoreIndexer(
+        connection_config=connection_config,
+        index_config=SingleStoreIndexerConfig(table_name="cars", id_column="car_id", batch_size=6),
+    )
+    downloader = SingleStoreDownloader(
+        connection_config=connection_config,
+        download_config=SingleStoreDownloaderConfig(
+            fields=["car_id", "brand"], download_dir=temp_dir
+        ),
+    )
+    await source_connector_validation(
+        indexer=indexer,
+        downloader=downloader,
+        configs=SourceValidationConfigs(
+            test_id="singlestore",
+            expected_num_files=SEED_DATA_ROWS,
+            expected_number_indexed_file_data=2,
+            validate_downloaded_files=True,
+        ),
+    )
+
+
+def validate_destination(
+    connect_params: dict,
+    expected_num_elements: int,
+):
+    with s2.connect(**connect_params) as connection:
+        with connection.cursor() as cursor:
+            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}"
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "sql")
+async def test_singlestore_destination(upload_file: Path, temp_dir: Path):
+    mock_file_data = FileData(
+        identifier="mock file data",
+        connector_type=CONNECTOR_TYPE,
+        source_identifiers=SourceIdentifiers(filename=upload_file.name, fullpath=upload_file.name),
+    )
+    with docker_compose_context(
+        docker_compose_path=env_setup_path / "sql" / "singlestore" / "destination"
+    ):
+        stager = SingleStoreUploadStager()
+        staged_path = stager.run(
+            elements_filepath=upload_file,
+            file_data=mock_file_data,
+            output_dir=temp_dir,
+            output_filename=upload_file.name,
+        )
+
+        # The stager should append the `.json` suffix to the output filename passed in.
+        assert staged_path.suffix == upload_file.suffix
+
+        connect_params = {
+            "host": "localhost",
+            "port": 3306,
+            "database": "ingest_test",
+            "user": "root",
+            "password": "password",
+        }
+
+        uploader = SingleStoreUploader(
+            connection_config=SingleStoreConnectionConfig(
+                host=connect_params["host"],
+                port=connect_params["port"],
+                database=connect_params["database"],
+                user=connect_params["user"],
+                access_config=SingleStoreAccessConfig(password=connect_params["password"]),
+            ),
+            upload_config=SingleStoreUploaderConfig(
+                table_name="elements",
+            ),
+        )
+        uploader.precheck()
+        uploader.run(path=staged_path, file_data=mock_file_data)
+
+        with staged_path.open("r") as f:
+            staged_data = json.load(f)
+        expected_num_elements = len(staged_data)
+        validate_destination(
+            connect_params=connect_params,
+            expected_num_elements=expected_num_elements,
+        )
+
+        uploader.run(path=staged_path, file_data=mock_file_data)
+        validate_destination(
+            connect_params=connect_params,
+            expected_num_elements=expected_num_elements,
+        )
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_singlestore_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = SingleStoreUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )

test/integration/connectors/sql/test_snowflake.py

@@ -0,0 +1,238 @@
+import json
+import os
+from pathlib import Path
+
+import pytest
+import snowflake.connector as sf
+from _pytest.fixtures import TopRequest
+
+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.destination import (
+    StagerValidationConfigs,
+    stager_validation,
+)
+from test.integration.connectors.utils.validation.source import (
+    SourceValidationConfigs,
+    source_connector_validation,
+)
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+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() -> dict:
+    connect_params = {
+        "user": "test",
+        "password": "test",
+        "account": "test",
+        "database": "test",
+        "host": "snowflake.localhost.localstack.cloud",
+    }
+    conn = sf.connect(**connect_params)
+
+    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()
+    return connect_params
+
+
+@pytest.fixture
+def source_database_setup() -> dict:
+    token = os.getenv("LOCALSTACK_AUTH_TOKEN")
+    with container_context(
+        image="localstack/snowflake",
+        environment={"LOCALSTACK_AUTH_TOKEN": token, "EXTRA_CORS_ALLOWED_ORIGINS": "*"},
+        ports={4566: 4566, 443: 443},
+        healthcheck_retries=30,
+    ):
+        connect_params = seed_data()
+        yield connect_params
+
+
+def init_db_destination() -> dict:
+    connect_params = {
+        "user": "test",
+        "password": "test",
+        "account": "test",
+        "database": "test",
+        "host": "snowflake.localhost.localstack.cloud",
+    }
+    conn = sf.connect(**connect_params)
+
+    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()
+    return connect_params
+
+
+@pytest.fixture
+def destination_database_setup() -> dict:
+    token = os.getenv("LOCALSTACK_AUTH_TOKEN")
+    with container_context(
+        image="localstack/snowflake",
+        environment={"LOCALSTACK_AUTH_TOKEN": token, "EXTRA_CORS_ALLOWED_ORIGINS": "*"},
+        ports={4566: 4566, 443: 443},
+        healthcheck_retries=30,
+    ):
+        connect_params = init_db_destination()
+        yield connect_params
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG, "sql")
+@requires_env("LOCALSTACK_AUTH_TOKEN")
+async def test_snowflake_source(temp_dir: Path, source_database_setup: dict):
+    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=temp_dir
+        ),
+    )
+    await source_connector_validation(
+        indexer=indexer,
+        downloader=downloader,
+        configs=SourceValidationConfigs(
+            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, temp_dir: Path, destination_database_setup: dict
+):
+    # 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,
+        source_identifiers=SourceIdentifiers(filename=upload_file.name, fullpath=upload_file.name),
+    )
+    init_db_destination()
+    stager = SnowflakeUploadStager()
+    staged_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=mock_file_data,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    # The stager should append the `.json` suffix to the output filename passed in.
+    assert staged_path.suffix == upload_file.suffix
+
+    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"],
+        )
+    )
+    uploader.precheck()
+    uploader.run(path=staged_path, file_data=mock_file_data)
+
+    with staged_path.open("r") as f:
+        staged_data = json.load(f)
+    expected_num_elements = len(staged_data)
+    validate_destination(
+        connect_params=connect_params,
+        expected_num_elements=expected_num_elements,
+    )
+
+    uploader.run(path=staged_path, file_data=mock_file_data)
+    validate_destination(
+        connect_params=connect_params,
+        expected_num_elements=expected_num_elements,
+    )
+
+
+@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
+def test_snowflake_stager(
+    request: TopRequest,
+    upload_file_str: str,
+    tmp_path: Path,
+):
+    upload_file: Path = request.getfixturevalue(upload_file_str)
+    stager = SnowflakeUploadStager()
+    stager_validation(
+        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
+        input_file=upload_file,
+        stager=stager,
+        tmp_dir=tmp_path,
+    )