"""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()