OpenMetadata/ingestion/tests/integration/cockroach/conftest.py

import os
import textwrap

import pytest
from sqlalchemy import create_engine

from _openmetadata_testutils.helpers.docker import try_bind
from metadata.generated.schema.api.services.createDatabaseService import (
    CreateDatabaseServiceRequest,
)
from metadata.generated.schema.entity.services.connections.database.cockroachConnection import (
    CockroachConnection,
)
from metadata.generated.schema.entity.services.databaseService import (
    DatabaseConnection,
    DatabaseServiceType,
)


@pytest.fixture(scope="module")
def cockroach_container(tmp_path_factory):
    """
    Start a Cockroach container.
    """
    from testcontainers.cockroachdb import CockroachDBContainer

    container = CockroachDBContainer(image="cockroachdb/cockroach:v23.1.0")

    with (
        try_bind(container, 26257, None) if not os.getenv("CI") else container
    ) as container:
        engine = create_engine(container.get_connection_url())
        engine.execute(
            textwrap.dedent(
                """
                CREATE TABLE user_profiles (
                    user_id UUID PRIMARY KEY,
                    first_name TEXT,
                    last_name TEXT,
                    email TEXT,
                    signup_date TIMESTAMP,
                    is_active BOOLEAN
                );
                """
            )
        )

        yield container


@pytest.fixture(scope="module")
def create_service_request(cockroach_container, tmp_path_factory):
    return CreateDatabaseServiceRequest(
        name="docker_test_" + tmp_path_factory.mktemp("cockroach").name,
        serviceType=DatabaseServiceType.Cockroach,
        connection=DatabaseConnection(
            config=CockroachConnection(
                username=cockroach_container.username,
                authType={"password": cockroach_container.password},
                hostPort=f"localhost:{cockroach_container.get_exposed_port(26257)}",
                database=cockroach_container.dbname,
            )
        ),
    )