From 1b230c38cd13829fe8729122ff40166b1b586299 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Wed, 26 Aug 2026 08:18:46 +0100 Subject: [PATCH 1/3] perf(api): stop replaying migrations, and give tests their own database The API suite took a little over twenty minutes, and a single test took around twenty-five seconds. Both were dominated by work that had nothing to do with the tests being run. Nearly all of it was migration replay. Creating a test database runs the whole ~600 migration history, and pytest-django does it once per xdist worker, so a cold ten-worker run paid for it ten times. Every migration test then replayed the history from zero to reach the state it tests, and migrated all the way forward again on teardown -- the 128 of them accounted for half the suite's CPU on their own. ClickHouse was the same story from a different angle: Django replayed all ~600 migrations against it to build the migration state for the three it owns, twenty-two seconds of it, on every session -- including a run of one Postgres-only test. PostgreSQL copies a database in about a tenth of a second, so cache each migration state as a template database and clone it instead. States are keyed on the length of the migration plan prefix they correspond to, which makes them nest: reaching a deeper state clones the nearest cached ancestor and replays only the migrations in between, so a session converges on the cost of one migration run however many migration tests it has. Template names embed a digest of the migration files, so editing a migration -- or switching branches -- builds a new template rather than handing anyone a stale schema. That is `--reuse-db` without the footgun, and it makes the `--ci` database bootstrap unnecessary. Test data is throwaway, which makes a Docker volume the wrong place for it: on macOS its I/O is virtualised, and it made the suite's runtime a function of how busy the host's disk was. Tests now run against `test-db`, a separate RAM-backed server tuned for a workload that clones databases and never needs to survive a crash. The rest is smaller change of the same character: the default password hasher costs 350ms a call and fixtures hash on almost every test; coverage roughly doubles the runtime and is now `make test-coverage`; `-vvvv` spent seconds formatting a line per test; and `make test` recreated every container and booted Django twice just to check the databases were up. Suite: ~1240s to ~150s on an idle machine. Individual runs: p95 6.7s over 40 sampled test files, against ~25s of fixed cost before a test even ran. --- .github/workflows/api-pull-request.yml | 4 +- api/.env-ci | 1 - api/.env-local | 6 +- api/Makefile | 24 +- api/README.md | 10 +- api/app/settings/test.py | 20 + api/pyproject.toml | 3 +- api/scripts/ensure-services.sh | 30 ++ api/tests/conftest.py | 110 +++--- api/tests/migration_snapshots.py | 489 +++++++++++++++++++++++++ docker/api/docker-compose.local.yml | 68 ++++ 11 files changed, 711 insertions(+), 54 deletions(-) create mode 100644 api/scripts/ensure-services.sh create mode 100644 api/tests/migration_snapshots.py diff --git a/.github/workflows/api-pull-request.yml b/.github/workflows/api-pull-request.yml index 21f446c9bb04..d084f3fad897 100644 --- a/.github/workflows/api-pull-request.yml +++ b/.github/workflows/api-pull-request.yml @@ -69,7 +69,9 @@ jobs: linters: mypy - name: Run Tests - run: make test + env: + format: xml + run: make test-coverage - name: Upload Coverage uses: codecov/codecov-action@v5 diff --git a/api/.env-ci b/api/.env-ci index efb12e4b6f8d..19d1fa9dd53b 100644 --- a/api/.env-ci +++ b/api/.env-ci @@ -5,5 +5,4 @@ CLICKHOUSE_PORT=9000 CLICKHOUSE_USER=flagsmith CLICKHOUSE_PASSWORD=password CLICKHOUSE_DATABASE=default -PYTEST_ADDOPTS=--cov . --cov-report xml -n auto --ci COVERAGE_CORE=sysmon diff --git a/api/.env-local b/api/.env-local index 28383d1233d2..62f7e92a31f0 100644 --- a/api/.env-local +++ b/api/.env-local @@ -5,5 +5,9 @@ CLICKHOUSE_PORT=9000 CLICKHOUSE_USER=flagsmith CLICKHOUSE_PASSWORD=password CLICKHOUSE_DATABASE=default +# The suite gets its own RAM-backed server, so it never competes with the +# database `make serve` keeps your local data in. See `test-db` in +# docker/api/docker-compose.local.yml. +TEST_DATABASE_URL=postgresql://postgres:password@localhost:5434/flagsmith +TEST_ANALYTICS_DATABASE_URL=postgresql://postgres:password@localhost:5434/analytics DJANGO_SETTINGS_MODULE=app.settings.local -PYTEST_ADDOPTS=--cov . --cov-report html -n auto diff --git a/api/Makefile b/api/Makefile index 374d7b938e97..695470f7b2d1 100644 --- a/api/Makefile +++ b/api/Makefile @@ -45,6 +45,14 @@ docker-up: docker compose up --force-recreate --remove-orphans -d docker compose ps +# `docker-up` recreates every container, which costs several seconds and drops +# open connections. Targets that only need the services to be reachable use +# this instead: it is near-free once the stack is up, and Compose's `--wait` +# blocks on the healthchecks so nothing has to poll for a database afterwards. +.PHONY: docker-ready +docker-ready: + @bash scripts/ensure-services.sh + .PHONY: docker-down docker-down: docker compose stop @@ -66,9 +74,21 @@ wait-for-db: uv run python manage.py waitfordb uv run python manage.py waitfordb --database analytics +# Every xdist worker pays for a Python start, a Django setup and a test +# database clone, so fanning out only earns its keep on a large run. A bare +# `make test` runs the whole suite and gets the fan-out; `make test opts=...` +# is a focused run and stays in-process. Put `-n auto` in `opts` to override. +PYTEST_PARALLELISM = $(if $(strip $(opts)),,-n auto) + .PHONY: test -test: docker-up wait-for-db - uv run pytest $(opts) +test: docker-ready + uv run pytest $(PYTEST_PARALLELISM) $(opts) + +# Coverage roughly doubles the suite's runtime, so it is opt-in rather than +# part of every `make test`. CI and anyone checking a diff's coverage use this. +.PHONY: test-coverage +test-coverage: docker-ready + uv run pytest --cov . --cov-report $(or $(format),html) $(PYTEST_PARALLELISM) $(opts) .PHONY: django-make-migrations django-make-migrations: docker-up wait-for-db diff --git a/api/README.md b/api/README.md index 0c2fb4d7ac7a..6cd33aca90a8 100644 --- a/api/README.md +++ b/api/README.md @@ -13,10 +13,18 @@ To run linters, run `make lint`. To run tests, run `make test`. -To run a subset of tests or an individual test, run `make test opts=''`. If the number of test is too low for xdist, consider adding `-n0` to pytest args. +To run a subset of tests or an individual test, run `make test opts=''`. A bare `make test` fans out across your cores with xdist; passing `opts` is taken to mean a focused run and stays in-process, since each worker costs a Python start, a Django setup and a database clone. Put `-n auto` in `opts` to fan out anyway. + +To measure coverage, run `make test-coverage`. It is not part of `make test` because it roughly doubles the runtime. To prepare a dev database, run `make docker-up django-migrate`. +#### Test databases + +Tests run against `test-db`, a separate PostgreSQL service that keeps its data in RAM. It is deliberately not the `db` service `make serve` uses: nothing the suite writes is worth keeping, and holding it in memory keeps the suite off the host's disk. Losing it costs nothing, so `docker compose restart test-db` is always a safe reset. + +You should never need `--reuse-db`. Rather than running the ~600 migrations for every worker on every run, the suite migrates once into a template database and clones it, which takes about a tenth of a second. The template is keyed on a digest of the migration files, so editing a migration -- or switching to a branch that has different ones -- silently builds a new template instead of handing you a stale schema. Migration tests use the same mechanism to reach the state they test, rather than replaying history for each one. See `tests/migration_snapshots.py`. + To bring up a dev server, run `make serve`, or `make serve-with-task-processor` to run the Task processor alongside the server. ### Code guidelines: testing diff --git a/api/app/settings/test.py b/api/app/settings/test.py index 72bd7920c906..468af6057ba7 100644 --- a/api/app/settings/test.py +++ b/api/app/settings/test.py @@ -1,15 +1,35 @@ +import dj_database_url + from app.settings.common import * # noqa from app.settings.common import ( DATABASES, INSTALLED_APPS, LDAP_INSTALLED, REST_FRAMEWORK, + env, ) +# Point the suite at its own PostgreSQL server when one is configured, so it +# never shares a disk -- or a page cache -- with the database `make serve` +# keeps your local data in. See `test-db` in the dev Compose file. +for alias, variable in ( + ("default", "TEST_DATABASE_URL"), + ("analytics", "TEST_ANALYTICS_DATABASE_URL"), +): + if alias in DATABASES and (url := env.str(variable, default="")): + DATABASES[alias] = {**DATABASES[alias], **dj_database_url.parse(url)} + +# Django's default hasher deliberately costs ~350ms per call. Fixtures like +# `admin_user` and `admin_master_api_key` hash a password on nearly every test, +# which made hashing one of the suite's largest line items. Nothing here +# depends on the algorithm, so use the cheapest one. +PASSWORD_HASHERS = ["django.contrib.auth.hashers.MD5PasswordHasher"] + # TODO: remove once permissions are an enum -- # https://github.com/Flagsmith/flagsmith/issues/7850 DATABASES["default"]["ENGINE"] = "core.db_backends.postgresql" + if LDAP_INSTALLED: INSTALLED_APPS = INSTALLED_APPS + ["flagsmith_ldap"] LDAP_DEFAULT_FLAGSMITH_ORGANISATION_ID = None diff --git a/api/pyproject.toml b/api/pyproject.toml index b9ebf22c75df..1dd6fa5e7000 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -475,7 +475,8 @@ omit = [ pythonpath = ['.'] addopts = [ '--ds=app.settings.test', - '-vvvv', + # `-v` and above print a line per test, which on a suite this size costs + # several seconds of pure formatting and buries failures in the noise. '-p', 'no:warnings', '--dist=worksteal', diff --git a/api/scripts/ensure-services.sh b/api/scripts/ensure-services.sh new file mode 100644 index 000000000000..0b3dc7d8937a --- /dev/null +++ b/api/scripts/ensure-services.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Make sure the local Compose services are reachable, cheaply. +# +# `docker compose up --wait` is the correct thing to run, but it needs a couple +# of seconds just to parse the project and talk to the daemon -- a large slice +# of a focused `make test`. Opening a socket to each service costs a +# millisecond and answers the same question in the overwhelmingly common case +# where the stack is already up, so try that first and fall back to Compose. +set -euo pipefail + +# Pull "host" and "port" out of a URL like postgresql://user:pw@host:5432/name. +endpoint_of() { + sed -E 's#^[^:]+://([^@]*@)?([^:/]+):([0-9]+).*$#\2 \3#' <<<"$1" +} + +reachable() { + # Bash's /dev/tcp needs no subprocess, unlike nc(1), which is not everywhere. + (exec 3<>"/dev/tcp/$1/$2") 2>/dev/null +} + +services=( + "$(endpoint_of "${TEST_DATABASE_URL:-${DATABASE_URL}}")" + "$(endpoint_of "${TEST_ANALYTICS_DATABASE_URL:-${ANALYTICS_DATABASE_URL}}")" + "${CLICKHOUSE_HOST:-localhost} ${CLICKHOUSE_PORT:-9000}" +) + +for service in "${services[@]}"; do + # shellcheck disable=SC2086 # deliberate split into host and port + reachable ${service} || exec docker compose up --remove-orphans --wait -d +done diff --git a/api/tests/conftest.py b/api/tests/conftest.py index 0ad59e5ccf92..950d3ea0557a 100644 --- a/api/tests/conftest.py +++ b/api/tests/conftest.py @@ -43,9 +43,7 @@ from django.conf import settings from django.contrib.contenttypes.models import ContentType from django.core.cache import caches -from django.db import connections -from django.db.backends.base.creation import TEST_DATABASE_PREFIX -from django.test.utils import setup_databases +from django.db import DEFAULT_DB_ALIAS, connections from django_test_migrations.migrator import Migrator from flag_engine.segments.constants import EQUAL from moto import mock_dynamodb # type: ignore[import-untyped] @@ -54,13 +52,12 @@ from pyfakefs.fake_filesystem import FakeFilesystem from pytest import FixtureRequest from pytest_django.fixtures import SettingsWrapper -from pytest_django.plugin import blocking_manager_key +from pytest_django.plugin import DjangoDbBlocker from pytest_mock import MockerFixture from rest_framework.test import APIClient from task_processor.task_run_method import TaskRunMethod from urllib3 import BaseHTTPResponse from urllib3.connectionpool import HTTPConnectionPool -from xdist import get_xdist_worker_id # type: ignore[import-untyped] from api_keys.models import MasterAPIKey from api_keys.user import APIKeyUser @@ -116,6 +113,12 @@ # TODO: Delete alias as per https://github.com/Flagsmith/flagsmith/issues/7818 from segments.types import SegmentRule as SegmentRuleType +from tests.migration_snapshots import ( + MigrationSnapshots, + SnapshotMigrator, + build_snapshots, + template_backed_test_databases, +) from tests.types import ( AdminClientAuthType, EnableFeaturesFixture, @@ -141,28 +144,15 @@ # --------------------------------------------------------------------------- -def pytest_addoption(parser: pytest.Parser) -> None: - parser.addoption( - "--ci", - action="store_true", - default=False, - help="Enable CI mode", - ) - +@pytest.fixture(scope="session", autouse=True) +def _template_backed_test_databases() -> typing.Generator[None, None, None]: + """Clone test databases from a migrated template rather than migrating. -@pytest.hookimpl(trylast=True) -def pytest_configure(config: pytest.Config) -> None: - if ( - config.option.ci - and config.option.dist != "no" - and not hasattr(config, "workerinput") - ): - with config.stash[blocking_manager_key].unblock(): - setup_databases( - verbosity=config.option.verbose, - interactive=False, - parallel=config.option.numprocesses, - ) + Autouse and session scoped so that it wraps `django_db_setup`, whichever + worker gets there first. See `tests.migration_snapshots`. + """ + with template_backed_test_databases(): + yield # --------------------------------------------------------------------------- @@ -213,27 +203,6 @@ def fs(fs: FakeFilesystem) -> FakeFilesystem: return fs -@pytest.fixture(scope="session") -def django_db_setup(request: pytest.FixtureRequest) -> None: - if ( - request.config.option.ci - # xdist worker id is either `gw[0-9]+` or `master` - and (xdist_worker_id_suffix := get_xdist_worker_id(request)[2:]).isnumeric() - ): - # Django's test database clone indices start at 1, - # Pytest's worker indices are 0-based - test_db_suffix = str(int(xdist_worker_id_suffix) + 1) - else: - # Tests are run on main node, which assumes -n0 - return request.getfixturevalue("django_db_setup") # type: ignore[no-any-return] # pragma: no cover - - from django.conf import settings - - for db_settings in settings.DATABASES.values(): - test_db_name = f"{TEST_DATABASE_PREFIX}{db_settings['NAME']}_{test_db_suffix}" - db_settings["NAME"] = test_db_name - - @pytest.fixture() def mock_influxdb_client(mocker: MockerFixture) -> MagicMock: client: MagicMock = mocker.patch.object(InfluxDBWrapper, "get_client").return_value @@ -1615,6 +1584,53 @@ def dynamo_environment_wrapper( return wrapper +@pytest.fixture(scope="session") +def migration_snapshots( + django_db_setup: None, + django_db_blocker: DjangoDbBlocker, +) -> typing.Generator[dict[str, MigrationSnapshots], None, None]: + """Cache each database's migration states as template databases. + + Built lazily, on the first migration test a worker runs, from a database + that `django_db_setup` has just migrated -- so the `latest` snapshot the + `migrator` fixture restores on teardown carries the rows that data + migrations create, which the old `migrate`-forward teardown could not. + """ + snapshots: dict[str, MigrationSnapshots] = {} + with django_db_blocker.unblock(): + yield snapshots + for snapshot in snapshots.values(): + snapshot.close() + + +@pytest.fixture() +def migrator_factory( + request: pytest.FixtureRequest, + transactional_db: None, + django_db_use_migrations: bool, + migration_snapshots: dict[str, MigrationSnapshots], +) -> MigratorFactory: + """Override `django_test_migrations`' fixture of the same name. + + Identical in behaviour -- including keeping the name, so the plugin still + recognises these tests and marks them `migration_test` -- except that + migration states are cloned from template databases rather than replayed + migration by migration. + """ + if not django_db_use_migrations: # pragma: no cover + pytest.skip("--no-migrations was specified") + + def factory(database_name: str | None = None) -> Migrator: + alias = database_name or DEFAULT_DB_ALIAS + if alias not in migration_snapshots: + migration_snapshots[alias] = build_snapshots(alias) + migrator = SnapshotMigrator(alias, migration_snapshots[alias]) + request.addfinalizer(migrator.reset) + return migrator + + return typing.cast(MigratorFactory, factory) + + @pytest.fixture() def migrator(migrator_factory: MigratorFactory) -> Migrator: if settings.SKIP_MIGRATION_TESTS: # pragma: no cover diff --git a/api/tests/migration_snapshots.py b/api/tests/migration_snapshots.py new file mode 100644 index 000000000000..963adf498ac4 --- /dev/null +++ b/api/tests/migration_snapshots.py @@ -0,0 +1,489 @@ +"""Snapshot-backed migration states, so tests stop replaying 500+ migrations. + +Two things in this suite pay for the migration history over and over: + +* Creating a test database runs every migration, and pytest-django does it + once per xdist worker -- so a cold ten-worker run pays for it ten times. +* Every migration test replays the history from zero up to the migration under + test, and then, on teardown, migrates all the way forward again. + +PostgreSQL can copy a whole database in about a tenth of a second with +`CREATE DATABASE ... TEMPLATE`, which is several hundred times faster than +replaying the history. So we cache migration states as template databases and +clone them. + +States are keyed on the length of the migration plan prefix they correspond +to. Django builds its "clean start" plan by walking the migration graph in a +deterministic order, so every state a test can ask for is a prefix of that one +plan, and the states nest: building the state for a prefix of length N clones +the deepest cached prefix shorter than N and replays only the migrations in +between. Over a session the cache converges on the cost of a single migration +run, however many migration tests there are. + +Template names embed a digest of the migration files on disk, so adding, +removing or editing a migration transparently invalidates every cached state. +That is what makes this safe to leave on by default, unlike `--reuse-db`: +there is no stale schema to notice and no flag to remember. Templates from +graphs that no longer exist are dropped when the cache is next used. +""" + +from __future__ import annotations + +import contextlib +import functools +import hashlib +import pathlib +import typing + +from django.apps import apps +from django.conf import settings as django_settings +from django.core.management.color import no_style +from django.db import DEFAULT_DB_ALIAS, connections +from django.db.backends.base.creation import BaseDatabaseCreation +from django.db.migrations.state import ProjectState +from django_test_migrations import sql +from django_test_migrations.logic.migrations import normalize +from django_test_migrations.migrator import Migrator +from django_test_migrations.plan import truncate_plan +from django_test_migrations.types import MigrationPlan, MigrationSpec + +# Length of the migration graph digest embedded in template database names. +# Eight hex characters comfortably separate the handful of graphs a working +# copy sees while keeping names well inside PostgreSQL's 63 byte limit. +_DIGEST_LENGTH = 8 + +# Shared by every template this module manages, so they are easy to recognise +# in `psql -l` and safe to drop wholesale. +_TEMPLATE_INFIX = "migsnap" + +# Key for the state in which the whole history has been applied. Named rather +# than numbered because it is the state the rest of the suite runs against. +_LATEST = "latest" + + +class SnapshotsUnavailable(Exception): + """Raised when the database cannot back migration state snapshots.""" + + +@functools.cache +def migration_graph_digest() -> str: + """Digest the migration files on disk. + + Hashes file contents rather than importing them through `MigrationLoader`: + it is an order of magnitude quicker, needs no database, and -- unlike + mtimes -- gives the same answer on a fresh clone as on a working copy, so + CI and a laptop agree on which templates they can share. + """ + digest = hashlib.sha256() + paths = sorted( + path + for app_config in apps.get_app_configs() + for path in pathlib.Path(app_config.path).glob("migrations/*.py") + ) + for path in paths: + digest.update(path.name.encode()) + digest.update(path.read_bytes()) + return digest.hexdigest()[:_DIGEST_LENGTH] + + +class _MaintenanceConnection: + """A connection to the maintenance database, for `CREATE`/`DROP DATABASE`. + + Neither statement may run inside a transaction block or touch the database + the issuing session is connected to, so they need a connection outside + Django's pool. + """ + + def __init__(self, alias: str) -> None: + self._alias = alias + self._connection: typing.Any = None + + def _connect(self) -> typing.Any: + if self._connection is None or self._connection.closed: + connection = connections[self._alias] + params = connection.get_connection_params() + params["dbname"] = "postgres" + params.pop("cursor_factory", None) + # `Database` is the DB-API module the backend was built against, + # which is how Django itself opens connections outside its pool. + driver = connection.Database # type: ignore[attr-defined] + self._connection = driver.connect(**params) + self._connection.set_session(autocommit=True) + return self._connection + + def execute(self, statement: str, params: typing.Sequence[typing.Any] = ()) -> None: + with self._connect().cursor() as cursor: + cursor.execute(statement, params) + + def fetch( + self, + statement: str, + params: typing.Sequence[typing.Any] = (), + ) -> list[tuple[typing.Any, ...]]: + with self._connect().cursor() as cursor: + cursor.execute(statement, params) + return list(cursor.fetchall()) + + def close(self) -> None: + if self._connection is not None and not self._connection.closed: + self._connection.close() + self._connection = None + + +@contextlib.contextmanager +def _advisory_lock( + maintenance: _MaintenanceConnection, + namespace: str, + *, + shared: bool, +) -> typing.Iterator[None]: + """Hold a PostgreSQL advisory lock naming `namespace`. + + Workers share the template databases, so cloning one must not overlap with + rebuilding it. Clones take the lock in shared mode and therefore run + concurrently with each other, which is by far the common case; builders + take it exclusively. + """ + mode = "_shared" if shared else "" + maintenance.execute(f"SELECT pg_advisory_lock{mode}(hashtext(%s))", (namespace,)) + try: + yield + finally: + maintenance.execute( + f"SELECT pg_advisory_unlock{mode}(hashtext(%s))", (namespace,) + ) + + +class MigrationSnapshots: + """Template databases holding migration states for one database alias.""" + + def __init__(self, alias: str, database_name: str | None = None) -> None: + connection = connections[alias] + if connection.vendor != "postgresql": # pragma: no cover + raise SnapshotsUnavailable( + f"Migration state snapshots need PostgreSQL, got {connection.vendor!r}" + ) + self._alias = alias + self._maintenance = _MaintenanceConnection(alias) + self._database_name = database_name or connection.settings_dict["NAME"] + self._namespace = self._build_namespace() + self._cached: set[str] = set() + self._discover() + + @property + def _base_name(self) -> str: + # Deliberately drops any xdist worker suffix: every worker migrates the + # same graph, so they should share the templates they build. + return self._database_name.split("_gw")[0] + + def _build_namespace(self) -> str: + return f"{self._base_name}_{_TEMPLATE_INFIX}_{migration_graph_digest()}" + + def _template_name(self, key: int | str) -> str: + return f"{self._namespace}_{key}" + + def _discover(self) -> None: + """Load the usable templates, dropping any left by an older graph. + + Scoped to this database's own templates: aliases can share a server -- + `default` and `analytics` do, on the dev stack's test server -- and one + alias must not mistake another's templates for its own leftovers. + """ + rows = self._maintenance.fetch( + "SELECT datname FROM pg_database WHERE datname LIKE %s", + (f"{self._base_name}\\_{_TEMPLATE_INFIX}\\_%",), + ) + for (name,) in rows: + if name.startswith(f"{self._namespace}_"): + self._cached.add(name) + else: + # Built from a migration graph that no longer exists. Leaving + # it would cost disk for every branch a working copy visits. + self._drop(name) + + def _drop(self, name: str) -> None: + # A template another worker is cloning right now cannot be dropped; + # whoever comes next will clean it up. + with contextlib.suppress(Exception): + self._maintenance.execute(f'DROP DATABASE IF EXISTS "{name}"') + + def _copy(self, source: str, target: str) -> None: + # PostgreSQL refuses to copy or drop a database that has sessions + # attached, and a worker's own connections are not always the only + # ones: a crashed run can leave backends behind. + self._maintenance.execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " + "WHERE datname IN (%s, %s) AND pid <> pg_backend_pid()", + (source, target), + ) + self._maintenance.execute(f'DROP DATABASE IF EXISTS "{target}"') + self._maintenance.execute(f'CREATE DATABASE "{target}" TEMPLATE "{source}"') + + def _advisory_lock(self, *, shared: bool) -> typing.ContextManager[None]: + return _advisory_lock(self._maintenance, self._namespace, shared=shared) + + def exclusive(self) -> typing.ContextManager[None]: + """Hold this namespace's lock for writing, to build a template. + + Without it, every xdist worker would migrate from scratch on a cold + cache instead of waiting a moment for the first one to finish. + """ + return self._advisory_lock(shared=False) + + def refresh(self) -> None: + """Re-read which templates exist, e.g. after waiting on the lock.""" + self._cached.clear() + self._discover() + + def has(self, key: int | str) -> bool: + return self._template_name(key) in self._cached + + def nearest_ancestor(self, depth: int) -> int | None: + """Return the deepest cached prefix strictly shorter than `depth`.""" + prefix = f"{self._namespace}_" + candidates = [ + int(suffix) + for name in self._cached + if (suffix := name[len(prefix) :]).isdigit() and int(suffix) < depth + ] + return max(candidates, default=None) + + def restore(self, key: int | str) -> None: + """Replace the working database with the cached state at `key`.""" + connections[self._alias].close() + with self._advisory_lock(shared=True): + self._copy(self._template_name(key), self._database_name) + + def save(self, key: int | str) -> None: + """Cache the working database's current state under `key`.""" + name = self._template_name(key) + connections[self._alias].close() + with self.exclusive(): + self._copy(self._database_name, name) + self._cached.add(name) + + def close(self) -> None: + self._maintenance.close() + + +class ClickHouseSnapshots: + """`CREATE DATABASE ... TEMPLATE` for ClickHouse, which has no such thing. + + ClickHouse only owns three of this project's migrations, but Django still + replays the whole history against the alias to build the migration state, + which is where nearly all of the twenty-odd seconds went. Copying the + handful of tables it does own -- schema via `CREATE TABLE ... AS`, plus the + `django_migrations` rows that say the history is already applied -- gets + the same database in well under a second. + """ + + # The one table whose *rows* matter: it is what stops Django replaying the + # migration history against the clone. + _MIGRATIONS_TABLE = "django_migrations" + + def __init__(self, alias: str, database_name: str | None = None) -> None: + self._alias = alias + self._database_name = database_name or connections[alias].settings_dict["NAME"] + base = self._database_name.split("_gw")[0] + self._template = f"{base}_{_TEMPLATE_INFIX}_{migration_graph_digest()}" + # ClickHouse has no advisory locks, so borrow PostgreSQL's. A mutex + # does not care which server it lives on, and every run that reaches + # here has the default database configured anyway. + self._maintenance = _MaintenanceConnection(DEFAULT_DB_ALIAS) + + def exclusive(self) -> typing.ContextManager[None]: + """Hold this template's lock for writing. See `MigrationSnapshots`.""" + return _advisory_lock(self._maintenance, self._template, shared=False) + + def shared(self) -> typing.ContextManager[None]: + """Hold this template's lock for reading, to clone it.""" + return _advisory_lock(self._maintenance, self._template, shared=True) + + def _quote(self, name: str) -> str: + return connections[self._alias].ops.quote_name(name) + + @contextlib.contextmanager + def _cursor(self) -> typing.Iterator[typing.Any]: + connection = connections[self._alias] + with connection._nodb_cursor() as cursor: # noqa: SLF001 + yield cursor + + def _databases(self) -> set[str]: + with self._cursor() as cursor: + cursor.execute("SELECT name FROM system.databases") + return {name for (name,) in cursor.fetchall()} + + def has(self, key: int | str) -> bool: + del key # ClickHouse only ever caches the fully migrated state. + return self._template in self._databases() + + def _copy(self, source: str, target: str) -> None: + with self._cursor() as cursor: + cursor.execute(f"DROP DATABASE IF EXISTS {self._quote(target)} SYNC") + cursor.execute(f"CREATE DATABASE {self._quote(target)}") + cursor.execute( + "SELECT name FROM system.tables WHERE database = %s", [source] + ) + for (table,) in cursor.fetchall(): + cursor.execute( + f"CREATE TABLE {self._quote(target)}.{self._quote(table)} " + f"AS {self._quote(source)}.{self._quote(table)}" + ) + migrations = self._quote(self._MIGRATIONS_TABLE) + cursor.execute( + f"INSERT INTO {self._quote(target)}.{migrations} " + f"SELECT * FROM {self._quote(source)}.{migrations}" + ) + + def restore(self, key: int | str) -> None: + del key + connections[self._alias].close() + with self.shared(): + self._copy(self._template, self._database_name) + + def save(self, key: int | str) -> None: + del key + connections[self._alias].close() + with self.exclusive(): + self._copy(self._database_name, self._template) + + def close(self) -> None: + self._maintenance.close() + + +class SnapshotMigrator(Migrator): + """A `Migrator` that clones cached states instead of replaying migrations.""" + + def __init__(self, database: str | None, snapshots: MigrationSnapshots) -> None: + super().__init__(database) + self._snapshots = snapshots + + def _full_plan(self) -> MigrationPlan: + self._executor.loader.build_graph() # reload + return self._executor.migration_plan( + self._executor.loader.graph.leaf_nodes(), + clean_start=True, + ) + + def apply_initial_migration(self, targets: MigrationSpec) -> ProjectState: + migration_targets = normalize(targets) + depth = len(truncate_plan(migration_targets, self._full_plan())) + + if self._snapshots.has(depth): + self._snapshots.restore(depth) + return self._restored_project_state() + + ancestor = self._snapshots.nearest_ancestor(depth) + if ancestor is None: + sql.drop_models_tables(self._database, no_style()) + sql.flush_django_migrations_table(self._database, no_style()) + ancestor = 0 + self._snapshots.save(ancestor) + else: + self._snapshots.restore(ancestor) + + # Replay only the migrations between the state we restored and the one + # under test. Rebuilding the plan first refreshes the executor's view + # of what the restored database has applied. + plan = self._full_plan()[ancestor:depth] + state = self._migrate(migration_targets, plan=plan) + self._snapshots.save(depth) + return state + + def reset(self) -> None: + """Restore the fully migrated state the rest of the suite expects.""" + self._snapshots.restore(_LATEST) + + def _restored_project_state(self) -> ProjectState: + """Build the model state matching the migrations the database records. + + Cloning a template leaves `django_migrations` exactly as the snapshot + had it, so the applied set is the source of truth and no migration has + to run to derive the historical models. + """ + self._executor.loader.build_graph() + state: ProjectState = self._executor._create_project_state( # type: ignore[attr-defined] + with_applied_migrations=True, + ) + state.clear_delayed_apps_cache() + return state + + +def build_snapshots(alias: str = DEFAULT_DB_ALIAS) -> MigrationSnapshots: + """Open the snapshot cache for `alias`, recording the migrated state.""" + snapshots = MigrationSnapshots(alias) + if not snapshots.has(_LATEST): + # Only reachable when the test database was not created through + # `template_backed_test_databases`, e.g. under `--reuse-db`. + with snapshots.exclusive(): # pragma: no cover + snapshots.refresh() + if not snapshots.has(_LATEST): + snapshots.save(_LATEST) + return snapshots + + +@contextlib.contextmanager +def _building( + snapshots: MigrationSnapshots | ClickHouseSnapshots, +) -> typing.Iterator[None]: + """Serialise template construction across xdist workers, where possible.""" + with snapshots.exclusive(): + if isinstance(snapshots, MigrationSnapshots): + snapshots.refresh() + yield + + +@contextlib.contextmanager +def template_backed_test_databases() -> typing.Iterator[None]: + """Make Django build test databases by cloning a migrated template. + + The first caller to arrive migrates and leaves a template behind; everyone + after that -- later xdist workers, later runs, other branches on the same + migration graph -- clones it. + """ + original = BaseDatabaseCreation.create_test_db + + def create_test_db( + self: BaseDatabaseCreation, + verbosity: int = 1, + autoclobber: bool = False, + serialize: bool = True, + keepdb: bool = False, + ) -> str: + alias = self.connection.alias + test_database_name: str = self._get_test_db_name() # type: ignore[attr-defined] + + snapshots: MigrationSnapshots | ClickHouseSnapshots + if self.connection.vendor == "clickhouse": + snapshots = ClickHouseSnapshots(alias, database_name=test_database_name) + else: + try: + snapshots = MigrationSnapshots(alias, database_name=test_database_name) + except SnapshotsUnavailable: # pragma: no cover + return original(self, verbosity, autoclobber, serialize, keepdb) + + try: + if not snapshots.has(_LATEST): + with _building(snapshots): + # Another worker may have built the template while we + # waited for the lock, in which case cloning it is still + # hundreds of times cheaper than migrating. + if not snapshots.has(_LATEST): + name = original(self, verbosity, autoclobber, serialize, keepdb) + snapshots.save(_LATEST) + return name + + snapshots.restore(_LATEST) + self.connection.close() + django_settings.DATABASES[alias]["NAME"] = test_database_name + self.connection.settings_dict["NAME"] = test_database_name + self.connection.ensure_connection() + return test_database_name + finally: + snapshots.close() + + BaseDatabaseCreation.create_test_db = create_test_db # type: ignore[method-assign] + try: + yield + finally: + BaseDatabaseCreation.create_test_db = original # type: ignore[method-assign] diff --git a/docker/api/docker-compose.local.yml b/docker/api/docker-compose.local.yml index eea96919d7c3..0cdbe8d2da2a 100644 --- a/docker/api/docker-compose.local.yml +++ b/docker/api/docker-compose.local.yml @@ -21,6 +21,11 @@ services: environment: POSTGRES_DB: flagsmith POSTGRES_PASSWORD: password + healthcheck: + test: ["CMD-SHELL", "pg_isready -U postgres -d flagsmith"] + interval: 2s + timeout: 3s + retries: 15 analytics-db: image: postgres:15.5-alpine @@ -34,6 +39,69 @@ services: environment: POSTGRES_DB: analytics POSTGRES_PASSWORD: password + healthcheck: + test: ["CMD-SHELL", "pg_isready -U postgres -d analytics"] + interval: 2s + timeout: 3s + retries: 15 + + # The test suite's own database, deliberately kept apart from `db` above. + # + # Nothing in here outlives a run: the suite clones a fresh database per xdist + # worker from a template it builds on demand, and drops them at the end. That + # makes durability pure overhead, and it makes a Docker volume -- whose I/O + # is virtualised, and slow, on macOS -- the suite's bottleneck. Holding it in + # RAM instead roughly halves the suite's wall time and, more importantly, + # stops the runtime depending on how busy the host's disk happens to be. + # + # Losing the lot on restart costs nothing. The template cache is keyed on a + # digest of the migrations, so an empty server simply rebuilds it. + # + # Both test databases live here: `default` and `analytics` are separate + # servers in production, and in `db`/`analytics-db` above, but tests only + # need them to be separate databases. + test-db: + image: postgres:15.5-alpine + pull_policy: always + restart: unless-stopped + command: + - "-c" + - "max_locks_per_transaction=256" + # Nothing here is worth surviving a crash. + - "-c" + - "fsync=off" + - "-c" + - "synchronous_commit=off" + - "-c" + - "full_page_writes=off" + # Cloning a database writes it through the WAL, so keep checkpoints out + # of the way of a run that clones one per worker and per migration test. + - "-c" + - "wal_level=minimal" + - "-c" + - "max_wal_senders=0" + - "-c" + - "max_wal_size=4GB" + - "-c" + - "checkpoint_timeout=30min" + # Comfortably fits every worker's database and the templates. + - "-c" + - "shared_buffers=1GB" + # One connection per worker, plus a maintenance connection each. + - "-c" + - "max_connections=200" + tmpfs: + - /var/lib/postgresql/data:size=3g + ports: + - 5434:5432 + environment: + POSTGRES_DB: flagsmith + POSTGRES_PASSWORD: password + healthcheck: + test: ["CMD-SHELL", "pg_isready -U postgres -d flagsmith"] + interval: 2s + timeout: 3s + retries: 15 influxdb: image: influxdb:2-alpine From e50dbd7d8c31bca45ff47f8547908f2ac6a3d58e Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Wed, 26 Aug 2026 10:58:13 +0100 Subject: [PATCH 2/3] fix(api): Skip the test databases' serialized_rollback snapshot Setting up the databases, Django serialises every model in every one of them so that `TransactionTestCase(serialized_rollback=True)` can restore them later. Nothing in this suite asks for that, so the work is pure cost. On the ClickHouse alias it is worse than pure cost. Serialising builds a `MigrationLoader`, and `django-clickhouse-backend` caches its migration model on `MigrationRecorder` -- which carries a `deleted` column that Django's does not. Get there with a PostgreSQL connection first and the cache holds a model without that column, while the queryset still filters on it, so the ClickHouse alias fails with FieldError: Cannot resolve keyword 'deleted' into field That only surfaces in the private build, where an app routed to ClickHouse gives the alias something to serialise. --- api/tests/migration_snapshots.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/api/tests/migration_snapshots.py b/api/tests/migration_snapshots.py index 963adf498ac4..f1f6de9fa461 100644 --- a/api/tests/migration_snapshots.py +++ b/api/tests/migration_snapshots.py @@ -482,8 +482,24 @@ def create_test_db( finally: snapshots.close() + def serialize_db_to_string(self: BaseDatabaseCreation) -> str: + """Skip the setup-time snapshot Django takes for `serialized_rollback`. + + Django serialises every model in every database while setting the + databases up, so that `TransactionTestCase(serialized_rollback=True)` + can restore them afterwards. Nothing in this suite asks for that, so + it is pure cost -- and on the ClickHouse alias it is worse than that: + it builds a `MigrationLoader`, and `django-clickhouse-backend` caches + its migration model on `MigrationRecorder` in a way that breaks if a + PostgreSQL connection got there first. + """ + return "" + + original_serialize = BaseDatabaseCreation.serialize_db_to_string BaseDatabaseCreation.create_test_db = create_test_db # type: ignore[method-assign] + BaseDatabaseCreation.serialize_db_to_string = serialize_db_to_string # type: ignore[method-assign] try: yield finally: BaseDatabaseCreation.create_test_db = original # type: ignore[method-assign] + BaseDatabaseCreation.serialize_db_to_string = original_serialize # type: ignore[method-assign] From a7061ce7830b8748c6d5f1020206e8244338a7c1 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Wed, 26 Aug 2026 15:51:31 +0100 Subject: [PATCH 3/3] test(api): Cover the migration snapshot cache's invalidation The cache drops templates whose migration graph no longer exists, which is what stops every branch a working copy visits leaving a copy of the database behind. CI never exercises it -- a job only ever sees one graph -- so assert it directly. --- api/tests/unit/test_migration_snapshots.py | 73 ++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 api/tests/unit/test_migration_snapshots.py diff --git a/api/tests/unit/test_migration_snapshots.py b/api/tests/unit/test_migration_snapshots.py new file mode 100644 index 000000000000..88c8cad1a145 --- /dev/null +++ b/api/tests/unit/test_migration_snapshots.py @@ -0,0 +1,73 @@ +import typing + +import pytest +from django.db import DEFAULT_DB_ALIAS, connections + +from tests.migration_snapshots import ( + MigrationSnapshots, + _MaintenanceConnection, + migration_graph_digest, +) + + +@pytest.fixture() +def maintenance(db: None) -> typing.Generator[_MaintenanceConnection, None, None]: + """A connection that can create and drop databases, as the cache uses.""" + connection = _MaintenanceConnection(DEFAULT_DB_ALIAS) + yield connection + connection.close() + + +@pytest.fixture() +def template_name_for_graph( + db: None, +) -> typing.Callable[[str], str]: + """Name a template as the cache would, for a given migration graph digest. + + Deliberately suffixed `probe` rather than `latest` or a plan depth: those + are the names the cache really uses, and a test that dropped one would pull + the database out from under every other xdist worker. + """ + + def name_for(digest: str) -> str: + base = connections[DEFAULT_DB_ALIAS].settings_dict["NAME"].split("_gw")[0] + return f"{base}_migsnap_{digest}_probe" + + return name_for + + +def test_migration_snapshots__template_from_another_graph__is_dropped( + db: None, + maintenance: _MaintenanceConnection, + template_name_for_graph: typing.Callable[[str], str], +) -> None: + """Templates only stay useful while their migrations do. + + A working copy that visits a branch with different migrations leaves a + template behind that can never be cloned again, so opening the cache drops + it. Without this, every branch a developer checks out would cost another + copy of the database. + """ + # Given + stale = template_name_for_graph("0ldgr4ph") + current = template_name_for_graph(migration_graph_digest()) + for name in (stale, current): + maintenance.execute(f'DROP DATABASE IF EXISTS "{name}"') + maintenance.execute(f'CREATE DATABASE "{name}"') + + # When + snapshots = MigrationSnapshots(DEFAULT_DB_ALIAS) + + # Then + try: + remaining = { + name + for (name,) in maintenance.fetch( + "SELECT datname FROM pg_database WHERE datname IN (%s, %s)", + (stale, current), + ) + } + assert remaining == {current} + finally: + snapshots.close() + maintenance.execute(f'DROP DATABASE IF EXISTS "{current}"')