unstructured-ingest 0.7.1__py3-none-any.whl → 1.0.1__py3-none-any.whl

test/integration/connectors/sql/test_snowflake.py

@@ -1,244 +0,0 @@
-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,
-    SQL_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.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.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_TAG)
-@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_TAG)
-@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.tags(CONNECTOR_TYPE, DESTINATION_TAG, SQL_TAG)
-@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,
-    )

test/integration/connectors/sql/test_sqlite.py

@@ -1,168 +0,0 @@
-import json
-import sqlite3
-import tempfile
-from pathlib import Path
-
-import pytest
-from _pytest.fixtures import TopRequest
-
-from test.integration.connectors.utils.constants import (
-    DESTINATION_TAG,
-    SOURCE_TAG,
-    SQL_TAG,
-    env_setup_path,
-)
-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.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.processes.connectors.sql.sqlite import (
-    CONNECTOR_TYPE,
-    SQLiteConnectionConfig,
-    SQLiteDownloader,
-    SQLiteDownloaderConfig,
-    SQLiteIndexer,
-    SQLiteIndexerConfig,
-    SQLiteUploader,
-    SQLiteUploadStager,
-)
-
-SEED_DATA_ROWS = 10
-
-
-@pytest.fixture
-def source_database_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_TAG)
-async def test_sqlite_source(source_database_setup: Path, temp_dir: Path):
-    connection_config = SQLiteConnectionConfig(database_path=source_database_setup)
-    indexer = SQLiteIndexer(
-        connection_config=connection_config,
-        index_config=SQLiteIndexerConfig(table_name="cars", id_column="car_id", batch_size=6),
-    )
-    downloader = SQLiteDownloader(
-        connection_config=connection_config,
-        download_config=SQLiteDownloaderConfig(fields=["car_id", "brand"], download_dir=temp_dir),
-    )
-    await source_connector_validation(
-        indexer=indexer,
-        downloader=downloader,
-        configs=SourceValidationConfigs(
-            test_id="sqlite",
-            expected_num_files=SEED_DATA_ROWS,
-            expected_number_indexed_file_data=2,
-            validate_downloaded_files=True,
-        ),
-    )
-
-
-@pytest.fixture
-def destination_database_setup(temp_dir: Path) -> 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.
-    db_path = temp_dir / "elements.db"
-    db_init_path = env_setup_path / "sql" / "sqlite" / "destination" / "sqlite-schema.sql"
-    assert db_init_path.exists()
-    assert db_init_path.is_file()
-    with sqlite3.connect(database=db_path) as sqlite_connection:
-        with db_init_path.open("r") as f:
-            query = f.read()
-        cursor = sqlite_connection.cursor()
-        cursor.executescript(query)
-    return db_path
-
-
-def validate_destination(db_path: Path, expected_num_elements: int):
-    # Run the following validations:
-    # * Check that the number of records in the table match the expected value
-    connection = None
-    try:
-        connection = sqlite3.connect(database=db_path)
-        query = "select count(*) from elements;"
-        cursor = connection.cursor()
-        cursor.execute(query)
-        count = cursor.fetchone()[0]
-        assert (
-            count == expected_num_elements
-        ), f"dest check failed: got {count}, expected {expected_num_elements}"
-    finally:
-        if connection:
-            connection.close()
-
-
-@pytest.mark.asyncio
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, SQL_TAG)
-async def test_sqlite_destination(
-    upload_file: Path, temp_dir: Path, destination_database_setup: 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,
-        source_identifiers=SourceIdentifiers(filename=upload_file.name, fullpath=upload_file.name),
-    )
-    stager = SQLiteUploadStager()
-    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
-
-    uploader = SQLiteUploader(
-        connection_config=SQLiteConnectionConfig(database_path=destination_database_setup)
-    )
-    uploader.precheck()
-    uploader.run(path=staged_path, file_data=mock_file_data)
-
-    with staged_path.open("r") as f:
-        staged_data = json.load(f)
-    validate_destination(db_path=destination_database_setup, expected_num_elements=len(staged_data))
-
-    uploader.run(path=staged_path, file_data=mock_file_data)
-    validate_destination(db_path=destination_database_setup, expected_num_elements=len(staged_data))
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, SQL_TAG)
-@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
-def test_sqlite_stager(
-    request: TopRequest,
-    upload_file_str: str,
-    tmp_path: Path,
-):
-    upload_file: Path = request.getfixturevalue(upload_file_str)
-    stager = SQLiteUploadStager()
-    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_vastdb.py

@@ -1,34 +0,0 @@
-from pathlib import Path
-
-import pytest
-from _pytest.fixtures import TopRequest
-
-from test.integration.connectors.utils.constants import DESTINATION_TAG, SQL_TAG
-from test.integration.connectors.utils.validation.destination import (
-    StagerValidationConfigs,
-    stager_validation,
-)
-from unstructured_ingest.processes.connectors.sql.vastdb import (
-    CONNECTOR_TYPE,
-    VastdbUploadStager,
-    VastdbUploadStagerConfig,
-)
-
-
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, SQL_TAG)
-@pytest.mark.parametrize("upload_file_str", ["upload_file_ndjson", "upload_file"])
-def test_vast_stager(
-    request: TopRequest,
-    upload_file_str: str,
-    tmp_path: Path,
-):
-    upload_file: Path = request.getfixturevalue(upload_file_str)
-    stager = VastdbUploadStager(
-        upload_stager_config=VastdbUploadStagerConfig(rename_columns_map={"page_number": "page"})
-    )
-    stager_validation(
-        configs=StagerValidationConfigs(test_id=CONNECTOR_TYPE, expected_count=22),
-        input_file=upload_file,
-        stager=stager,
-        tmp_dir=tmp_path,
-    )