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..f1f6de9fa461 --- /dev/null +++ b/api/tests/migration_snapshots.py @@ -0,0 +1,505 @@ +"""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() + + 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] 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}"') 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