DataOps-platform/tests/integration/test_datasource_pools.py

"""Opt-in acceptance tests against isolated PostgreSQL and MySQL sources."""

import os
import threading
from concurrent.futures import ThreadPoolExecutor

import pytest
from sqlalchemy import text
from sqlalchemy.engine import make_url

from app.config.config import datasource_pool_settings
from app.core.data_source.adapters import adapter_for
from app.core.data_source.errors import (
    DataSourceNotFound,
    DataSourceReadOnlyViolation,
)
from app.core.data_source.manager import DataSourceConnectionManager
from app.core.data_source.models import (
    DataSourceCredential,
    DataSourceDefinition,
)
from app.core.data_source.pool_registry import PoolRegistry


SOURCE_URL_ENVIRONMENTS = (
    ("postgresql", "TEST_SOURCE_POSTGRES_URL"),
    ("mysql", "TEST_SOURCE_MYSQL_URL"),
)

pytestmark = pytest.mark.integration


class DefinitionRepository:
    def __init__(self, definition):
        self._definitions = {definition.uid: definition}

    def get(self, uid):
        return self._definitions.get(uid)

    def delete(self, uid):
        self._definitions.pop(uid, None)


class CredentialRepository:
    def __init__(self, credential):
        self.credential = credential

    def get_active(self, _session, _uid, _version):
        return self.credential


class FakeClock:
    def __init__(self):
        self._value = 1000.0
        self._lock = threading.Lock()

    def __call__(self):
        with self._lock:
            return self._value

    def advance(self, seconds):
        with self._lock:
            self._value += seconds


def _source_url(database_type, environment_name):
    raw = os.environ.get(environment_name)
    if not raw:
        pytest.skip(
            f"set {environment_name} to run the isolated {database_type} test"
        )
    return make_url(raw)


def _manager(database_type, environment_name):
    url = _source_url(database_type, environment_name)
    uid = (
        "01900000-0000-7000-8000-000000000101"
        if database_type == "postgresql"
        else "01900000-0000-7000-8000-000000000102"
    )
    definition = DataSourceDefinition(
        uid=uid,
        name_en=f"acceptance-{database_type}",
        database_type=database_type,
        host=url.host,
        port=url.port,
        database=url.database,
        credential_ref=uid,
        credential_version=1,
        pool_size=2,
        max_overflow=3,
    )
    credential = DataSourceCredential(
        username=url.username,
        password=url.password,
    )
    definitions = DefinitionRepository(definition)
    clock = FakeClock()
    settings = datasource_pool_settings()
    registry = PoolRegistry(settings, clock=clock)
    manager = DataSourceConnectionManager(
        definitions=definitions,
        credentials=CredentialRepository(credential),
        platform_session=lambda: object(),
        adapter_resolver=adapter_for,
        registry=registry,
        settings_resolver=datasource_pool_settings,
        clock=clock,
    )
    return uid, definitions, clock, registry, manager


@pytest.mark.parametrize(
    ("database_type", "environment_name"),
    SOURCE_URL_ENVIRONMENTS,
)
def test_real_source_pool_is_bounded_reused_evicted_and_recreated(
    database_type,
    environment_name,
):
    uid, definitions, clock, registry, manager = _manager(
        database_type,
        environment_name,
    )

    def read_customers(_index):
        with manager.connect(uid, purpose="dataflow_read") as connection:
            return tuple(
                connection.execute(
                    text(
                        "SELECT customer_name "
                        "FROM acceptance_customers ORDER BY id"
                    )
                ).scalars()
            )

    with ThreadPoolExecutor(max_workers=25) as executor:
        results = list(executor.map(read_customers, range(100)))

    assert set(results) == {("Alpha", "Beta")}
    status = registry.snapshot()[0]
    assert status.pool_size == 2
    assert status.checked_out == 0
    assert 1 < status.checked_out_peak <= 5
    assert status.query_total == 100

    with pytest.raises(DataSourceReadOnlyViolation):
        with manager.connect(uid, purpose="dataflow_read") as connection:
            connection.execute(
                text(
                    "INSERT INTO acceptance_customers "
                    "(id, customer_name) VALUES (999, 'Forbidden')"
                )
            )
    assert read_customers(100) == ("Alpha", "Beta")

    clock.advance(901)
    assert registry.evict_idle() == 1
    assert registry.snapshot() == []

    assert read_customers(101) == ("Alpha", "Beta")
    recreated = registry.snapshot()[0]
    assert recreated.connection_created_total == 1

    definitions.delete(uid)
    manager.invalidate(uid, "source deleted")
    with pytest.raises(DataSourceNotFound):
        with manager.connect(uid, purpose="dataflow_read"):
            pass
    manager.close()