unstructured-ingest 0.7.2__py3-none-any.whl → 1.0.2__py3-none-any.whl

test/integration/connectors/sql/test_databricks_delta_tables.py

@@ -1,170 +0,0 @@
-import json
-import os
-import time
-from contextlib import contextmanager
-from pathlib import Path
-from uuid import uuid4
-
-import pytest
-from databricks.sql import connect
-from databricks.sql.client import Connection as DeltaTableConnection
-from databricks.sql.client import Cursor as DeltaTableCursor
-from pydantic import BaseModel, Secret, SecretStr
-from pytest_mock import MockerFixture
-
-from test.integration.connectors.utils.constants import DESTINATION_TAG, SQL_TAG, env_setup_path
-from test.integration.utils import requires_env
-from unstructured_ingest.data_types.file_data import FileData, SourceIdentifiers
-from unstructured_ingest.logger import logger
-from unstructured_ingest.processes.connectors.sql.databricks_delta_tables import (
-    CONNECTOR_TYPE,
-    DatabricksDeltaTablesAccessConfig,
-    DatabricksDeltaTablesConnectionConfig,
-    DatabricksDeltaTablesUploader,
-    DatabricksDeltaTablesUploaderConfig,
-    DatabricksDeltaTablesUploadStager,
-)
-
-CATALOG = "utic-dev-tech-fixtures"
-
-
-class EnvData(BaseModel):
-    server_hostname: str
-    http_path: str
-    access_token: SecretStr
-
-
-def get_env_data() -> EnvData:
-    return EnvData(
-        server_hostname=os.environ["DATABRICKS_SERVER_HOSTNAME"],
-        http_path=os.environ["DATABRICKS_HTTP_PATH"],
-        access_token=os.environ["DATABRICKS_ACCESS_TOKEN"],
-    )
-
-
-def get_destination_schema(new_table_name: str) -> str:
-    p = Path(env_setup_path / "sql" / "databricks_delta_tables" / "destination" / "schema.sql")
-    with p.open() as f:
-        data_lines = f.readlines()
-    data_lines[0] = data_lines[0].replace("elements", new_table_name)
-    data = "".join([line.strip() for line in data_lines])
-    return data
-
-
-@contextmanager
-def get_connection() -> DeltaTableConnection:
-    env_data = get_env_data()
-    with connect(
-        server_hostname=env_data.server_hostname,
-        http_path=env_data.http_path,
-        access_token=env_data.access_token.get_secret_value(),
-    ) as connection:
-        yield connection
-
-
-@contextmanager
-def get_cursor() -> DeltaTableCursor:
-    with get_connection() as connection:
-        with connection.cursor() as cursor:
-            cursor.execute(f"USE CATALOG '{CATALOG}'")
-            yield cursor
-
-
-@pytest.fixture
-def destination_table() -> str:
-    random_id = str(uuid4())[:8]
-    table_name = f"elements_{random_id}"
-    destination_schema = get_destination_schema(new_table_name=table_name)
-    with get_cursor() as cursor:
-        logger.info(f"creating table: {table_name}")
-        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
-        cursor.execute(destination_schema)
-
-    yield table_name
-    with get_cursor() as cursor:
-        logger.info(f"dropping table: {table_name}")
-        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
-
-
-def validate_destination(expected_num_elements: int, table_name: str, retries=30, interval=1):
-    with get_cursor() as cursor:
-        for i in range(retries):
-            cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
-            count = cursor.fetchone()[0]
-            if count == expected_num_elements:
-                break
-            logger.info(f"retry attempt {i}: expected {expected_num_elements} != count {count}")
-            time.sleep(interval)
-        assert (
-            count == expected_num_elements
-        ), f"dest check failed: got {count}, expected {expected_num_elements}"
-
-
-@pytest.mark.asyncio
-@pytest.mark.skip("Resources take too long to spin up to run in CI")
-@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, SQL_TAG)
-@requires_env("DATABRICKS_SERVER_HOSTNAME", "DATABRICKS_HTTP_PATH", "DATABRICKS_ACCESS_TOKEN")
-async def test_databricks_delta_tables_destination(
-    upload_file: Path, temp_dir: Path, destination_table: str
-):
-    env_data = get_env_data()
-    mock_file_data = FileData(
-        identifier="mock file data",
-        connector_type=CONNECTOR_TYPE,
-        source_identifiers=SourceIdentifiers(filename=upload_file.name, fullpath=upload_file.name),
-    )
-    stager = DatabricksDeltaTablesUploadStager()
-    staged_path = stager.run(
-        elements_filepath=upload_file,
-        file_data=mock_file_data,
-        output_dir=temp_dir,
-        output_filename=upload_file.name,
-    )
-
-    assert staged_path.suffix == upload_file.suffix
-
-    uploader = DatabricksDeltaTablesUploader(
-        connection_config=DatabricksDeltaTablesConnectionConfig(
-            access_config=DatabricksDeltaTablesAccessConfig(
-                token=env_data.access_token.get_secret_value()
-            ),
-            http_path=env_data.http_path,
-            server_hostname=env_data.server_hostname,
-        ),
-        upload_config=DatabricksDeltaTablesUploaderConfig(
-            catalog=CATALOG, database="default", table_name=destination_table
-        ),
-    )
-    with staged_path.open("r") as f:
-        staged_data = json.load(f)
-    expected_num_elements = len(staged_data)
-    uploader.precheck()
-    uploader.run(path=staged_path, file_data=mock_file_data)
-    validate_destination(expected_num_elements=expected_num_elements, table_name=destination_table)
-
-
-def test_get_credentials_provider_with_client_id_and_secret(mocker: MockerFixture):
-    access_config = DatabricksDeltaTablesAccessConfig(
-        client_id="test_client_id", client_secret="test_client_secret"
-    )
-    connection_config = DatabricksDeltaTablesConnectionConfig(
-        access_config=Secret(access_config),
-        server_hostname="test_server_hostname",
-        http_path="test_http_path",
-    )
-
-    credentials_provider = connection_config.get_credentials_provider()
-    assert credentials_provider is not False
-    assert type(credentials_provider).__name__ == "function"
-
-
-def test_get_credentials_provider_with_token(mocker: MockerFixture):
-    access_config = DatabricksDeltaTablesAccessConfig(token="test_token")
-    connection_config = DatabricksDeltaTablesConnectionConfig(
-        access_config=Secret(access_config),
-        server_hostname="test_server_hostname",
-        http_path="test_http_path",
-    )
-
-    credentials_provider = connection_config.get_credentials_provider()
-    assert credentials_provider is False

test/integration/connectors/sql/test_postgres.py

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

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