From dd8a33f03ef047d9654ccfdf2616a795d4359509 Mon Sep 17 00:00:00 2001
From: Prakhar Pratyush <prakhar@zulip.com>
Date: Tue, 12 Dec 2023 03:56:38 +0530
Subject: [PATCH] import_realm: Create audit log with user count data.

This commit creates a RealmAuditlog entry with a new event_type
'RealmAuditLog.REALM_IMPORTED' after the realm is reactivated.

It contains user count data (using realm_user_count_by_role)
stored in extra_data.

This helps to have an accurate user count data for the billing
system if someone tries to signup just after doing an import.
---
 zerver/lib/import_realm.py          | 12 ++++++++++++
 zerver/models.py                    |  2 ++
 zerver/tests/test_import_export.py  |  6 +++++-
 zerver/tests/test_slack_importer.py |  1 +
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py
index fbb219dd5d..7ceaf49bcb 100644
--- a/zerver/lib/import_realm.py
+++ b/zerver/lib/import_realm.py
@@ -35,6 +35,7 @@ from zerver.lib.timestamp import datetime_to_timestamp
 from zerver.lib.upload import upload_backend
 from zerver.lib.upload.base import BadImageError, sanitize_name
 from zerver.lib.upload.s3 import get_bucket
+from zerver.lib.user_counts import realm_user_count_by_role
 from zerver.lib.user_groups import create_system_user_groups_for_realm
 from zerver.lib.user_message import UserMessageLite, bulk_insert_ums
 from zerver.lib.utils import generate_api_key, process_list_in_batches
@@ -1419,6 +1420,17 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea
     realm.deactivated = data["zerver_realm"][0]["deactivated"]
     realm.save()
 
+    # This helps to have an accurate user count data for the billing
+    # system if someone tries to signup just after doing import.
+    RealmAuditLog.objects.create(
+        realm=realm,
+        event_type=RealmAuditLog.REALM_IMPORTED,
+        event_time=timezone_now(),
+        extra_data={
+            RealmAuditLog.ROLE_COUNT: realm_user_count_by_role(realm),
+        },
+    )
+
     # Ask the push notifications service if this realm can send
     # notifications, if we're using it. Needs to happen after the
     # Realm object is reactivated.
diff --git a/zerver/models.py b/zerver/models.py
index 1e7d227520..529965cba8 100644
--- a/zerver/models.py
+++ b/zerver/models.py
@@ -4802,6 +4802,7 @@ class AbstractRealmAuditLog(models.Model):
     REALM_EMOJI_ADDED = 226
     REALM_EMOJI_REMOVED = 227
     REALM_LINKIFIERS_REORDERED = 228
+    REALM_IMPORTED = 229
 
     SUBSCRIPTION_CREATED = 301
     SUBSCRIPTION_ACTIVATED = 302
@@ -4873,6 +4874,7 @@ class AbstractRealmAuditLog(models.Model):
         USER_ROLE_CHANGED,
         REALM_DEACTIVATED,
         REALM_REACTIVATED,
+        REALM_IMPORTED,
     ]
 
     class Meta:
diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py
index 5594e7063a..b36ecfdd29 100644
--- a/zerver/tests/test_import_export.py
+++ b/zerver/tests/test_import_export.py
@@ -1137,7 +1137,11 @@ class RealmImportExportTest(ExportFile):
         @getter
         def get_realm_audit_log_event_type(r: Realm) -> Set[int]:
             realmauditlogs = RealmAuditLog.objects.filter(realm=r).exclude(
-                event_type__in=[RealmAuditLog.REALM_PLAN_TYPE_CHANGED, RealmAuditLog.STREAM_CREATED]
+                event_type__in=[
+                    RealmAuditLog.REALM_PLAN_TYPE_CHANGED,
+                    RealmAuditLog.STREAM_CREATED,
+                    RealmAuditLog.REALM_IMPORTED,
+                ]
             )
             realmauditlog_event_type = {log.event_type for log in realmauditlogs}
             return realmauditlog_event_type
diff --git a/zerver/tests/test_slack_importer.py b/zerver/tests/test_slack_importer.py
index 1f22f9822d..b1c6d8473a 100644
--- a/zerver/tests/test_slack_importer.py
+++ b/zerver/tests/test_slack_importer.py
@@ -1336,6 +1336,7 @@ class SlackImporter(ZulipTestCase):
                 RealmAuditLog.SUBSCRIPTION_CREATED,
                 RealmAuditLog.REALM_PLAN_TYPE_CHANGED,
                 RealmAuditLog.REALM_CREATED,
+                RealmAuditLog.REALM_IMPORTED,
                 RealmAuditLog.USER_GROUP_CREATED,
                 RealmAuditLog.USER_GROUP_DIRECT_USER_MEMBERSHIP_ADDED,
                 RealmAuditLog.USER_GROUP_DIRECT_SUBGROUP_MEMBERSHIP_ADDED,