snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

zerver/lib/test_fixtures.py

@@ -0,0 +1,409 @@
+import glob
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+
+from django.conf import settings
+from django.db import DEFAULT_DB_ALIAS, ProgrammingError, connection, connections
+from django.db.utils import OperationalError
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+from scripts.lib.zulip_tools import (
+    TEMPLATE_DATABASE_DIR,
+    get_dev_uuid_var_path,
+    is_digest_obsolete,
+    run,
+    write_new_digest,
+)
+
+BACKEND_DATABASE_TEMPLATE = "zulip_test_template"
+UUID_VAR_DIR = get_dev_uuid_var_path()
+
+IMPORTANT_FILES = [
+    "zilencer/management/commands/populate_db.py",
+    "zerver/actions/create_realm.py",
+    "zerver/lib/bulk_create.py",
+    "zerver/lib/generate_test_data.py",
+    "zerver/lib/server_initialization.py",
+    "tools/setup/postgresql-init-test-db",
+    "tools/setup/postgresql-init-dev-db",
+    "zerver/migrations/0258_enable_online_push_notifications_default.py",
+]
+
+VERBOSE_MESSAGE_ABOUT_HASH_TRANSITION = """
+    NOTE!!!!
+
+    We are rebuilding your database for a one-time transition.
+
+    We have a hashing scheme that we use to detect whether any
+    important files used in the construction of the database
+    have changed.
+
+    We are changing that scheme so it only uses one file
+    instead of a directory of files.
+
+    In order to prevent errors due to this transition, we are
+    doing a one-time rebuild of your database.  This should
+    be the last time this happens (for this particular reason,
+    at least), unless you go back to older branches.
+
+"""
+
+
+def migration_paths() -> list[str]:
+    return [
+        *glob.glob("*/migrations/*.py"),
+        "uv.lock",
+    ]
+
+
+class Database:
+    def __init__(self, platform: str, database_name: str, settings: str) -> None:
+        self.database_name = database_name
+        self.settings = settings
+        self.digest_name = "db_files_hash_for_" + platform
+        self.migration_status_file = "migration_status_" + platform
+        self.migration_status_path = os.path.join(
+            UUID_VAR_DIR,
+            self.migration_status_file,
+        )
+        self.migration_digest_file = "migrations_hash_" + database_name
+
+    def important_settings(self) -> list[str]:
+        def get(setting_name: str) -> str:
+            value = getattr(settings, setting_name, {})
+            return json.dumps(value, sort_keys=True)
+
+        return [
+            get("LOCAL_DATABASE_PASSWORD"),
+            get("INTERNAL_BOTS"),
+            get("REALM_INTERNAL_BOTS"),
+            get("DISABLED_REALM_INTERNAL_BOTS"),
+        ]
+
+    def run_db_migrations(self) -> None:
+        # We shell out to `manage.py` and pass `DJANGO_SETTINGS_MODULE` on
+        # the command line rather than just calling the migration
+        # functions, because Django doesn't support changing settings like
+        # what the database is as runtime.
+        # Also we export ZULIP_DB_NAME which is ignored by dev platform but
+        # recognised by test platform and used to migrate correct db.
+        manage_py = [
+            "env",
+            "DJANGO_SETTINGS_MODULE=" + self.settings,
+            "ZULIP_DB_NAME=" + self.database_name,
+            "./manage.py",
+        ]
+
+        run([*manage_py, "migrate", "--skip-checks", "--no-input"])
+
+        run(
+            [
+                *manage_py,
+                "get_migration_status",
+                "--skip-checks",
+                "--output=" + self.migration_status_file,
+            ]
+        )
+
+    def what_to_do_with_migrations(self) -> str:
+        from zerver.lib.migration_status import get_migration_status
+
+        status_fn = self.migration_status_path
+        settings = self.settings
+
+        if not os.path.exists(status_fn):
+            return "scrap"
+
+        with open(status_fn) as f:
+            previous_migration_status = f.read()
+
+        current_migration_status = get_migration_status(settings=settings)
+        connections.close_all()
+        all_curr_migrations = extract_migrations_as_list(current_migration_status)
+        all_prev_migrations = extract_migrations_as_list(previous_migration_status)
+
+        if len(all_curr_migrations) < len(all_prev_migrations):
+            return "scrap"
+
+        for migration in all_prev_migrations:
+            if migration not in all_curr_migrations:
+                return "scrap"
+
+        if len(all_curr_migrations) == len(all_prev_migrations):
+            return "migrations_are_latest"
+
+        return "migrate"
+
+    def database_exists(self) -> bool:
+        try:
+            connection = connections[DEFAULT_DB_ALIAS]
+
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    "SELECT 1 from pg_database WHERE datname=%s;",
+                    [self.database_name],
+                )
+                return_value = bool(cursor.fetchone())
+            connections.close_all()
+            return return_value
+        except OperationalError:
+            return False
+
+    def files_or_settings_have_changed(self) -> bool:
+        database_name = self.database_name
+
+        # Deal with legacy hash files.  We can kill off this code when
+        # enough time has passed since April 2020 that we're not
+        # worried about anomalies doing `git bisect`--probably a few
+        # months is sufficient.
+        legacy_status_dir = os.path.join(UUID_VAR_DIR, database_name + "_db_status")
+        if os.path.exists(legacy_status_dir):
+            print(VERBOSE_MESSAGE_ABOUT_HASH_TRANSITION)
+
+            # Remove the old digest for several reasons:
+            #   - tidiness
+            #   - preventing false positives if you bisect
+            #   - make this only a one-time headache (generally)
+            shutil.rmtree(legacy_status_dir)
+
+            # Return True to force a one-time rebuild.
+            return True
+
+        return is_digest_obsolete(
+            self.digest_name,
+            IMPORTANT_FILES,
+            self.important_settings(),
+        )
+
+    def template_status(self) -> str:
+        # This function returns a status string specifying the type of
+        # state the template db is in and thus the kind of action required.
+        if not self.database_exists():
+            # TODO: It's possible that `database_exists` will
+            #       return `False` even though the database
+            #       exists, but we just have the wrong password,
+            #       probably due to changing the secrets file.
+            #
+            #       The only problem this causes is that we waste
+            #       some time rebuilding the whole database, but
+            #       it's better to err on that side, generally.
+            return "needs_rebuild"
+
+        if self.files_or_settings_have_changed():
+            return "needs_rebuild"
+
+        # Here we hash and compare our migration files before doing
+        # the work of seeing what to do with them; if there are no
+        # changes, we can safely assume we don't need to run
+        # migrations without spending a few 100ms parsing all the
+        # Python migration code.
+        if not self.is_migration_digest_obsolete():
+            return "current"
+
+        """
+        NOTE:
+            We immediately update the digest, assuming our
+            callers will do what it takes to run the migrations.
+
+            Ideally our callers would just do it themselves
+            AFTER the migrations actually succeeded, but the
+            caller codepaths are kind of complicated here.
+        """
+        self.write_new_migration_digest()
+
+        migration_op = self.what_to_do_with_migrations()
+        if migration_op == "scrap":
+            return "needs_rebuild"
+
+        if migration_op == "migrate":
+            return "run_migrations"
+
+        return "current"
+
+    def is_migration_digest_obsolete(self) -> bool:
+        return is_digest_obsolete(
+            self.migration_digest_file,
+            migration_paths(),
+        )
+
+    def write_new_migration_digest(self) -> None:
+        write_new_digest(
+            self.migration_digest_file,
+            migration_paths(),
+        )
+
+    def write_new_db_digest(self) -> None:
+        write_new_digest(
+            self.digest_name,
+            IMPORTANT_FILES,
+            self.important_settings(),
+        )
+
+
+DEV_DATABASE = Database(
+    platform="dev",
+    database_name="zulip",
+    settings="zproject.settings",
+)
+
+TEST_DATABASE = Database(
+    platform="test",
+    database_name="zulip_test_template",
+    settings="zproject.test_settings",
+)
+
+
+def update_test_databases_if_required(rebuild_test_database: bool = False) -> None:
+    """Checks whether the zulip_test_template database template, is
+    consistent with our database migrations; if not, it updates it
+    in the fastest way possible:
+
+    * If all we need to do is add some migrations, just runs those
+      migrations on the template database.
+    * Otherwise, we rebuild the test template database from scratch.
+
+    The default behavior is sufficient for the `test-backend` use
+    case, where the test runner code will clone directly from the
+    template database.
+
+    The `rebuild_test_database` option (used by our frontend and API
+    tests) asks us to drop and re-cloning the zulip_test database from
+    the template so those test suites can run with a fresh copy.
+
+    """
+    test_template_db_status = TEST_DATABASE.template_status()
+
+    if test_template_db_status == "needs_rebuild":
+        run(["tools/rebuild-test-database"])
+        TEST_DATABASE.write_new_db_digest()
+        return
+
+    if test_template_db_status == "run_migrations":
+        TEST_DATABASE.run_db_migrations()
+        run(["tools/setup/generate-fixtures"])
+        return
+
+    if rebuild_test_database:
+        run(["tools/setup/generate-fixtures"])
+
+
+def extract_migrations_as_list(migration_status: str) -> list[str]:
+    MIGRATIONS_RE = re.compile(r"\[[X| ]\] (\d+_.+)\n")
+    return MIGRATIONS_RE.findall(migration_status)
+
+
+def destroy_leaked_test_databases(expiry_time: int = 60 * 60) -> int:
+    """The logic in zerver/lib/test_runner.py tries to delete all the
+    temporary test databases generated by test-backend threads, but it
+    cannot guarantee it handles all race conditions correctly.  This
+    is a catch-all function designed to delete any that might have
+    been leaked due to crashes (etc.).  The high-level algorithm is to:
+
+    * Delete every database with a name like zulip_test_template_*
+    * Unless it is registered in a file under TEMPLATE_DATABASE_DIR as
+      part of a currently running test-backend invocation
+    * And that file is less expiry_time old.
+
+    This should ensure we ~never break a running test-backend process,
+    while also ensuring we will eventually delete all leaked databases.
+    """
+    files = glob.glob(os.path.join(UUID_VAR_DIR, TEMPLATE_DATABASE_DIR, "*"))
+    test_databases: set[str] = set()
+    try:
+        with connection.cursor() as cursor:
+            cursor.execute("SELECT datname FROM pg_database;")
+            rows = cursor.fetchall()
+            for row in rows:
+                if "zulip_test_template_" in row[0]:
+                    test_databases.add(row[0])
+    except ProgrammingError:
+        pass
+
+    databases_in_use: set[str] = set()
+    for file in files:
+        if round(time.time()) - os.path.getmtime(file) < expiry_time:
+            with open(file) as f:
+                databases_in_use.update(f"zulip_test_template_{line}".rstrip() for line in f)
+        else:
+            # Any test-backend run older than expiry_time can be
+            # cleaned up, both the database and the file listing its
+            # databases.
+            os.remove(file)
+
+    databases_to_drop = test_databases - databases_in_use
+
+    if not databases_to_drop:
+        return 0
+
+    commands = "\n".join(f"DROP DATABASE IF EXISTS {db};" for db in databases_to_drop)
+    subprocess.run(
+        ["psql", "-q", "-v", "ON_ERROR_STOP=1", "-h", "localhost", "postgres", "zulip_test"],
+        input=commands,
+        check=True,
+        text=True,
+    )
+    return len(databases_to_drop)
+
+
+def remove_test_run_directories(expiry_time: int = 60 * 60) -> int:
+    removed = 0
+    directories = glob.glob(os.path.join(UUID_VAR_DIR, "test-backend", "run_*"))
+    for test_run in directories:
+        if round(time.time()) - os.path.getmtime(test_run) > expiry_time:
+            try:
+                shutil.rmtree(test_run)
+                removed += 1
+            except FileNotFoundError:
+                pass
+    return removed
+
+
+def reset_zulip_test_database() -> None:
+    """
+    This function is used to reset the zulip_test database fastest way possible,
+    i.e. First, it deletes the database and then clones it from zulip_test_template.
+    This function is used with puppeteer tests, so it can quickly reset the test
+    database after each run.
+    """
+    from zerver.lib.test_runner import destroy_test_databases
+
+    # Make sure default database is 'zulip_test'.
+    assert connections["default"].settings_dict["NAME"] == "zulip_test"
+
+    # Clearing all the active PSQL sessions with 'zulip_test'.
+    run(
+        [
+            "env",
+            "PGHOST=localhost",
+            "PGUSER=zulip_test",
+            "scripts/setup/terminate-psql-sessions",
+            "zulip_test",
+        ]
+    )
+
+    destroy_test_databases()
+    # Pointing default database to test database template, so we can instantly clone it.
+    settings.DATABASES["default"]["NAME"] = BACKEND_DATABASE_TEMPLATE
+    connection = connections["default"]
+    clone_database_suffix = "clone"
+    connection.creation.clone_test_db(
+        suffix=clone_database_suffix,
+    )
+    settings_dict = connection.creation.get_test_db_clone_settings(clone_database_suffix)
+    # We manually rename the clone database to 'zulip_test' because when cloning it,
+    # its name is set to original database name + some suffix.
+    # Also, we need it to be 'zulip_test' so that our running server can recognize it.
+    with connection.cursor() as cursor:
+        cursor.execute("ALTER DATABASE zulip_test_template_clone RENAME TO zulip_test;")
+    settings_dict["NAME"] = "zulip_test"
+    # connection.settings_dict must be updated in place for changes to be
+    # reflected in django.db.connections. If the following line assigned
+    # connection.settings_dict = settings_dict, new threads would connect
+    # to the default database instead of the appropriate clone.
+    connection.settings_dict.update(settings_dict)
+    connection.close()