open-metadata--openmetadata/ingestion/tests/integration/cockroach/conftest.py

import os
import textwrap
import uuid

import pytest
from sqlalchemy import create_engine, text

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():
    """
    Start a Cockroach container.
    """
    from testcontainers.cockroachdb import CockroachDBContainer
    from testcontainers.core.config import testcontainers_config

    old_max_tries = testcontainers_config.max_tries
    testcontainers_config.max_tries = 240

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

    with try_bind(container, 26257, None) if not os.getenv("CI") else container as container:
        testcontainers_config.max_tries = old_max_tries
        engine = create_engine(container.get_connection_url())
        with engine.connect() as conn:
            conn.execute(
                text(
                    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
                    );
                    """
                    )
                )
            )
            conn.commit()

        yield container


@pytest.fixture(scope="module")
def create_service_request(cockroach_container):
    return CreateDatabaseServiceRequest(
        name=f"docker_test_cockroach_{uuid.uuid4().hex[:8]}",
        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,
            )
        ),
    )


@pytest.fixture(scope="module")
def create_test_data(cockroach_container):
    engine = create_engine(cockroach_container.get_connection_url())

    setup_statements = [
        """
        INSERT INTO user_profiles (user_id, first_name, last_name, email, signup_date, is_active)
        VALUES (gen_random_uuid(), 'Alice', 'Smith', 'alice.smith@example.com', '2023-01-15 10:00:00', TRUE)
        """,
        """
        INSERT INTO user_profiles (user_id, first_name, last_name, email, signup_date, is_active)
        VALUES (gen_random_uuid(), 'Bob', 'Jones', 'bob.jones@example.com', '2023-02-20 12:30:00', TRUE)
        """,
        """
        INSERT INTO user_profiles (user_id, first_name, last_name, email, signup_date, is_active)
        VALUES (gen_random_uuid(), 'Carol', 'Williams', 'carol.williams@example.com', '2023-03-10 09:15:00', FALSE)
        """,
        """
        INSERT INTO user_profiles (user_id, first_name, last_name, email, signup_date, is_active)
        VALUES (gen_random_uuid(), 'David', 'Brown', 'david.brown@example.com', '2023-04-05 14:45:00', TRUE)
        """,
        """
        INSERT INTO user_profiles (user_id, first_name, last_name, email, signup_date, is_active)
        VALUES (gen_random_uuid(), 'Eve', 'Davis', 'eve.davis@example.com', '2023-05-22 08:00:00', FALSE)
        """,
        """
        CREATE TABLE kv (
            k INT8 NOT NULL,
            v BYTES NOT NULL,
            CONSTRAINT kv_pkey PRIMARY KEY (k ASC)
        )
        """,
        "INSERT INTO kv (k, v) VALUES (1, b'\\x00\\x01\\x02\\x03')",
        "INSERT INTO kv (k, v) VALUES (2, b'\\x68\\x65\\x6c\\x6c\\x6f')",
        "INSERT INTO kv (k, v) VALUES (3, b'\\xde\\xad\\xbe\\xef')",
        "INSERT INTO kv (k, v) VALUES (4, b'\\xca\\xfe\\xba\\xbe')",
        "INSERT INTO kv (k, v) VALUES (5, b'\\xff\\x00\\xff\\x00')",
        """
        CREATE TABLE employees (
            employee_id INT8 PRIMARY KEY DEFAULT unique_rowid(),
            full_name TEXT NOT NULL,
            email TEXT NOT NULL,
            phone TEXT,
            department TEXT,
            hire_date TIMESTAMP
        )
        """,
        """
        INSERT INTO employees (full_name, email, phone, department, hire_date)
        VALUES ('John Doe', 'john.doe@company.com', '555-0101', 'Engineering', '2021-06-01 09:00:00')
        """,
        """
        INSERT INTO employees (full_name, email, phone, department, hire_date)
        VALUES ('Jane Smith', 'jane.smith@company.com', '555-0102', 'Marketing', '2021-07-15 09:00:00')
        """,
        """
        INSERT INTO employees (full_name, email, phone, department, hire_date)
        VALUES ('Robert Johnson', 'robert.johnson@company.com', '555-0103', 'Finance', '2022-01-10 09:00:00')
        """,
        """
        INSERT INTO employees (full_name, email, phone, department, hire_date)
        VALUES ('Maria Garcia', 'maria.garcia@company.com', '555-0104', 'HR', '2022-03-20 09:00:00')
        """,
        """
        INSERT INTO employees (full_name, email, phone, department, hire_date)
        VALUES ('James Wilson', 'james.wilson@company.com', '555-0105', 'Engineering', '2023-02-14 09:00:00')
        """,
    ]

    with engine.connect() as conn:
        for stmt in setup_statements:
            conn.execute(text(textwrap.dedent(stmt)))
        conn.commit()

    yield