import secrets
import time
from pathlib import Path
import sqlalchemy
import sqlalchemy.exc
from packaging.version import parse as version_parse
from requests.exceptions import ConnectionError # noqa A004
from sqlalchemy.dialects import registry
from sqlalchemy.orm import sessionmaker
from trino.exceptions import TrinoQueryError
from ehrql.query_engines.in_memory_database import InMemoryDatabase
from ehrql.utils.itertools_utils import iter_flatten
from tests.lib.orm_utils import SYNTHETIC_PRIMARY_KEY, table_has_one_row_per_patient
MSSQL_SETUP_DIR = Path(__file__).parents[1].absolute() / "support/mssql"
TRINO_SETUP_DIR = Path(__file__).parents[1].absolute() / "support/trino"
# Register our modified SQLAlchemy dialects
registry.register(
"sqlite.pysqlite.opensafely",
"ehrql.query_engines.sqlite_dialect",
"SQLiteDialect",
)
registry.register(
"trino.opensafely", "ehrql.query_engines.trino_dialect", "TrinoDialect"
)
class DbDetails:
def __init__(
self,
protocol,
driver,
host_from_container,
port_from_container,
host_from_host,
port_from_host,
username="",
password="",
db_name="",
query=None,
temp_db=None,
engine_kwargs=None,
):
self.protocol = protocol
self.driver = driver
self.host_from_container = host_from_container
self.port_from_container = port_from_container
self.host_from_host = host_from_host
self.port_from_host = port_from_host
self.password = password
self.username = username
self.db_name = db_name
self.query = query
self.temp_db = temp_db
self.engine_kwargs = engine_kwargs or {}
self.metadata = None
def container_url(self):
return self._url(self.host_from_container, self.port_from_container)
def host_url(self):
return self._url(self.host_from_host, self.port_from_host)
def engine(self, dialect=None, **kwargs):
url = self._url(
self.host_from_host, self.port_from_host, include_driver=bool(self.driver)
)
engine_url = sqlalchemy.engine.make_url(url)
engine_kwargs = self.engine_kwargs | kwargs
engine = sqlalchemy.create_engine(engine_url, **engine_kwargs)
return engine
def _url(self, host, port, include_driver=False):
assert self.username
if self.username and self.password:
auth = f"{self.username}:{self.password}@"
else:
auth = f"{self.username}@"
if include_driver:
protocol = f"{self.protocol}+{self.driver}"
else:
protocol = self.protocol
url = f"{protocol}://{auth}{host}:{port}/{self.db_name}"
return url
def setup(self, *input_data, metadata=None):
"""
Accepts SQLAlchemy ORM objects (which may be arbitrarily nested within lists and
tuples), creates the necessary tables and inserts them into the database
"""
input_data = list(iter_flatten(input_data))
engine = self.engine()
Session = sessionmaker()
Session.configure(bind=engine)
session = Session()
if metadata:
pass
elif input_data:
metadata = input_data[0].metadata
else:
assert False, "No source of metadata"
assert all(item.metadata is metadata for item in input_data)
self.metadata = metadata
metadata.create_all(engine)
session.bulk_save_objects(
input_data,
return_defaults=False,
update_changed_only=False,
preserve_order=False,
)
session.commit()
def teardown(self):
if self.metadata is not None:
self.metadata.drop_all(self.engine())
def wait_for_database(database, timeout=20):
engine = database.engine()
start = time.time()
limit = start + timeout
while True:
try:
with engine.connect() as connection:
connection.execute(sqlalchemy.text("SELECT 'hello'"))
break
except (
sqlalchemy.exc.OperationalError,
ConnectionRefusedError,
ConnectionResetError,
BrokenPipeError,
ConnectionError,
TrinoQueryError,
sqlalchemy.exc.DBAPIError,
) as e: # pragma: no cover
if time.time() >= limit:
raise Exception(
f"Failed to connect to database after {timeout} seconds: "
f"{engine.url}"
) from e
time.sleep(1)
def make_mssql_database(containers):
password = "Your_password123!"
container_name = "ehrql-mssql"
mssql_port = 1433
if not containers.is_running(container_name): # pragma: no cover
run_mssql(container_name, containers, password, mssql_port)
container_ip = containers.get_container_ip(container_name)
host_mssql_port = containers.get_mapped_port_for_host(container_name, mssql_port)
return DbDetails(
protocol="mssql",
driver="pymssql",
host_from_container=container_ip,
port_from_container=mssql_port,
host_from_host="localhost",
port_from_host=host_mssql_port,
username="sa",
password=password,
db_name="test",
)
def run_mssql(container_name, containers, password, mssql_port): # pragma: no cover
containers.run_bg(
name=container_name,
# This is *not* the version that TPP run for us in production which, as of
# 2024-09-24, is SQL Server 2016 (13.0.5893.48). That version is not available
# as a Docker image, so we run the oldest supported version instead. Both the
# production server and our test server set the "compatibility level" to the
# same value so the same feature set should be supported.
image="mcr.microsoft.com/mssql/server:2019-CU28-ubuntu-20.04",
volumes={
MSSQL_SETUP_DIR: {"bind": "/mssql", "mode": "ro"},
},
# Choose an arbitrary free port to publish the MSSQL port on
ports={mssql_port: None},
environment={
"MSSQL_SA_PASSWORD": password,
"ACCEPT_EULA": "Y",
"MSSQL_TCP_PORT": str(mssql_port),
# Make all string comparisons case-sensitive across all databases
"MSSQL_COLLATION": "SQL_Latin1_General_CP1_CS_AS",
},
user="root",
entrypoint="/mssql/entrypoint.sh",
command="/opt/mssql/bin/sqlservr",
)
class InMemorySQLiteDatabase(DbDetails):
def __init__(self):
db_name = secrets.token_hex(8)
super().__init__(
db_name=db_name,
protocol="sqlite",
driver="pysqlite+opensafely",
host_from_container=None,
port_from_container=None,
host_from_host=None,
port_from_host=None,
)
self._engine = None
def engine(self, dialect=None, **kwargs):
# We need to hold a reference to the engine for the lifetime of this database to stop the contents of the
# database from being garbage-collected.
if not self._engine:
self._engine = super().engine(dialect, **kwargs)
return self._engine
def _url(self, host, port, include_driver=False):
if include_driver:
protocol = f"{self.protocol}+{self.driver}"
else:
protocol = self.protocol
# https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#uri-connections
# https://sqlite.org/inmemorydb.html
return f"{protocol}:///file:{self.db_name}?mode=memory&cache=shared&uri=true"
class InMemoryPythonDatabase:
def __init__(self):
self.database = InMemoryDatabase()
def setup(self, *input_data, metadata=None):
"""
Behaves like `DbDetails.setup` in taking a iterator of ORM instances but
translates these into the sort of objects needed by the `InMemoryDatabase`
"""
input_data = list(iter_flatten(input_data))
if metadata:
pass
elif input_data:
metadata = input_data[0].metadata
else:
assert False, "No source of metadata"
assert all(item.metadata is metadata for item in input_data)
sqla_table_to_items = {table: [] for table in metadata.sorted_tables}
for item in input_data:
sqla_table_to_items[item.__table__].append(item)
for sqla_table, items in sqla_table_to_items.items():
columns = [
c.name for c in sqla_table.columns if c.name != SYNTHETIC_PRIMARY_KEY
]
self.database.add_table(
name=sqla_table.name,
one_row_per_patient=table_has_one_row_per_patient(sqla_table),
columns=columns,
rows=[[getattr(item, c) for c in columns] for item in items],
)
def teardown(self):
self.database.populate({})
def host_url(self):
# Where other query engines expect a DSN string to connect to the database the
# InMemoryQueryEngine expects a reference to the database object itself
return self.database
def make_trino_database(containers):
container_name = "ehrql-trino"
trino_port = 8080
if not containers.is_running(container_name): # pragma: no cover
run_trino(container_name, containers, trino_port)
container_ip = containers.get_container_ip(container_name)
host_trino_port = containers.get_mapped_port_for_host(container_name, trino_port)
return DbDetails(
protocol="trino",
driver="opensafely",
host_from_container=container_ip,
port_from_container=trino_port,
host_from_host="localhost",
port_from_host=host_trino_port,
username="trino",
db_name="trino/default",
# Disable automatic retries for the test client: it's pointless and creates log
# noise
engine_kwargs={"connect_args": {"max_attempts": 1}},
)
def run_trino(container_name, containers, trino_port): # pragma: no cover
# Note, I don't actually know that this is the minimum required version of Docker
# Engine. I do know that 20.10.5 is unsupported (because that's what I had
# installed) and that 20.10.16 is supported, according to this comment:
# https://github.com/adoptium/containers/issues/214#issuecomment-1139464798 which
# was linked from this issue: https://github.com/trinodb/trino/issues/14269
min_docker_version = "20.10.16"
docker_version = containers.get_engine_version()
assert version_parse(docker_version) >= version_parse(min_docker_version), (
f"The Trino Docker image requires Docker Engine v{min_docker_version}"
f" or above but you have v{docker_version}"
)
containers.run_bg(
name=container_name,
# This is the version which happened to be current at the time of writing and is
# pinned for reproduciblity's sake rather than because there's anything
# significant about it
image="trinodb/trino:440",
volumes={
TRINO_SETUP_DIR: {"bind": "/trino", "mode": "ro"},
f"{TRINO_SETUP_DIR}/etc": {"bind": "/etc/trino", "mode": "ro"},
},
# Choose an arbitrary free port to publish the trino port on
ports={trino_port: None},
environment={},
user="root",
entrypoint="/trino/entrypoint.sh",
command="/usr/lib/trino/bin/run-trino",
)
