From 40bcb4b42bbe762ff2789611fa3aef04a60019cc Mon Sep 17 00:00:00 2001 From: PieterCK Date: Tue, 29 Oct 2024 16:00:57 +0700 Subject: [PATCH] export: Add migration status file to export tarball. This commit updates the export process to write the migration status of the realm as a JSON file to be included in the export tarball. This is a preparatory step for adding an assertion to ensure that the importing and exporting realms have a compatible set of applied migrations. --- zerver/lib/export.py | 40 +++++++++++ .../import_fixtures/migration_status.json | 68 +++++++++++++++++++ zerver/tests/test_import_export.py | 24 ++++++- 3 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 zerver/tests/fixtures/import_fixtures/migration_status.json diff --git a/zerver/lib/export.py b/zerver/lib/export.py index 568a856ab2..107ce0f5a4 100644 --- a/zerver/lib/export.py +++ b/zerver/lib/export.py @@ -31,6 +31,7 @@ from django.utils.timezone import now as timezone_now import zerver.lib.upload from analytics.models import RealmCount, StreamCount, UserCount from scripts.lib.zulip_tools import overwrite_symlink +from version import ZULIP_VERSION from zerver.lib.avatar_hash import user_avatar_base_path_from_ids from zerver.lib.pysa import mark_sanitized from zerver.lib.timestamp import datetime_to_timestamp @@ -97,6 +98,8 @@ SourceFilter: TypeAlias = Callable[[Record], bool] CustomFetch: TypeAlias = Callable[[TableData, Context], None] +AppMigrations: TypeAlias = dict[str, list[str]] + class MessagePartial(TypedDict): zerver_message: list[Record] @@ -104,6 +107,11 @@ class MessagePartial(TypedDict): realm_id: int +class MigrationStatusJson(TypedDict): + migrations_by_app: AppMigrations + zulip_version: str + + MESSAGE_BATCH_CHUNK_SIZE = 1000 ALL_ZULIP_TABLES = { @@ -2093,6 +2101,8 @@ def do_export_realm( export_full_with_consent=export_type == RealmExport.EXPORT_FULL_WITH_CONSENT, ) + do_common_export_processes(output_dir) + logging.info("Finished exporting %s", realm.string_id) create_soft_link(source=output_dir, in_progress=False) @@ -2594,3 +2604,33 @@ def get_realm_exports_serialized(realm: Realm) -> list[dict[str, Any]]: export_type=export.type, ) return sorted(exports_dict.values(), key=lambda export_dict: export_dict["id"]) + + +def get_migrations_by_app() -> AppMigrations: + from django.db import DEFAULT_DB_ALIAS, connections + from django.db.migrations.recorder import MigrationRecorder + + recorder = MigrationRecorder(connections[DEFAULT_DB_ALIAS]) + applied = recorder.applied_migrations() + migrations_by_app: AppMigrations = {} + for app_name, migration_name in applied: + migrations_by_app.setdefault(app_name, []).append(migration_name) + return migrations_by_app + + +def export_migration_status(output_dir: str) -> None: + migration_status_json = MigrationStatusJson( + migrations_by_app=get_migrations_by_app(), zulip_version=ZULIP_VERSION + ) + output_file = os.path.join(output_dir, "migration_status.json") + with open(output_file, "wb") as f: + f.write(orjson.dumps(migration_status_json, option=orjson.OPT_INDENT_2)) + + +def do_common_export_processes(output_dir: str) -> None: + # Performs common task(s) necessary for preparing Zulip data exports. + # This function is typically shared with migration tools in the + # `zerver/data_import` directory. + + logging.info("Exporting migration status") + export_migration_status(output_dir) diff --git a/zerver/tests/fixtures/import_fixtures/migration_status.json b/zerver/tests/fixtures/import_fixtures/migration_status.json new file mode 100644 index 0000000000..06064d97e7 --- /dev/null +++ b/zerver/tests/fixtures/import_fixtures/migration_status.json @@ -0,0 +1,68 @@ +{ + "migrations_by_app": { + "contenttypes": ["0001_initial", "0002_remove_content_type_name"], + "auth": [ + "0001_initial", + "0002_alter_permission_name_max_length", + "0003_alter_user_email_max_length", + "0004_alter_user_username_opts" + ], + "zerver": [ + "0001_initial", + "0029_realm_subdomain", + "0030_realm_org_type", + "0031_remove_system_avatar_source" + ], + "analytics": [ + "0001_initial", + "0002_remove_huddlecount", + "0003_fillstate" + ], + "confirmation": [ + "0001_initial", + "0002_realmcreationkey", + "0003_emailchangeconfirmation" + ], + "zilencer": [ + "0001_initial", + "0002_remote_zulip_server", + "0003_add_default_for_remotezulipserver_last_updated_field" + ], + "corporate": [ + "0001_initial", + "0002_customer_default_discount", + "0003_customerplan" + ], + "otp_static": ["0001_initial", "0002_throttling", "0003_add_timestamps"], + "otp_totp": ["0001_initial", "0002_auto_20190420_0723", "0003_add_timestamps"], + "pgroonga": ["0001_enable", "0002_html_escape_subject", "0003_v2_api_upgrade"], + "phonenumber": ["0001_initial", "0001_squashed_0001_initial"], + "two_factor": [ + "0001_initial", + "0002_auto_20150110_0810", + "0003_auto_20150817_1733", + "0004_auto_20160205_1827" + ], + "sessions": ["0001_initial"], + "default": [ + "0001_initial", + "0002_add_related_name", + "0003_alter_email_max_length", + "0004_auto_20160423_0400" + ], + "social_auth": [ + "0001_initial", + "0002_add_related_name", + "0003_alter_email_max_length", + "0004_auto_20160423_0400", + "0005_auto_20160727_2333" + ], + "social_django": [ + "0006_partial", + "0007_code_timestamp", + "0008_partial_timestamp", + "0009_auto_20191118_0520" + ] + }, + "zulip_version": "10.0-dev+git" + } diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index cfda844d94..3b71d30e9d 100644 --- a/zerver/tests/test_import_export.py +++ b/zerver/tests/test_import_export.py @@ -42,7 +42,13 @@ from zerver.lib import upload from zerver.lib.avatar_hash import user_avatar_path from zerver.lib.bot_config import set_bot_config from zerver.lib.bot_lib import StateHandler -from zerver.lib.export import Record, do_export_realm, do_export_user, export_usermessages_batch +from zerver.lib.export import ( + MigrationStatusJson, + Record, + do_export_realm, + do_export_user, + export_usermessages_batch, +) from zerver.lib.import_realm import do_import_realm, get_incoming_message_ids from zerver.lib.streams import create_stream_if_needed from zerver.lib.test_classes import ZulipTestCase @@ -334,6 +340,20 @@ class ExportFile(ZulipTestCase): db_paths = {user_avatar_path(user) + ".original"} self.assertEqual(exported_paths, db_paths) + def verify_migration_status_json(self) -> None: + # This function asserts that the generated migration_status.json + # is structurally familiar for it to be used for assertion at + # import_realm.py. Hence, it doesn't really matter if the individual + # apps' migrations in migration_status.json fixture are outdated. + exported: MigrationStatusJson = read_json("migration_status.json") + fixture: MigrationStatusJson = orjson.loads( + self.fixture_data("migration_status.json", "import_fixtures") + ) + for app, migrations in fixture["migrations_by_app"].items(): + self.assertTrue( + set(migrations).issubset(set(exported["migrations_by_app"].get(app, []))), + ) + class RealmImportExportTest(ExportFile): def create_user_and_login(self, email: str, realm: Realm) -> None: @@ -392,6 +412,7 @@ class RealmImportExportTest(ExportFile): self.verify_avatars(user) self.verify_emojis(user, is_s3=False) self.verify_realm_logo_and_icon() + self.verify_migration_status_json() def test_public_only_export_files_private_uploads_not_included(self) -> None: """ @@ -439,6 +460,7 @@ class RealmImportExportTest(ExportFile): self.verify_avatars(user) self.verify_emojis(user, is_s3=True) self.verify_realm_logo_and_icon() + self.verify_migration_status_json() def test_zulip_realm(self) -> None: realm = Realm.objects.get(string_id="zulip")