From 0f36abb9865ed3e9989d0d480a026b1b047d0311 Mon Sep 17 00:00:00 2001
From: Steve Astels <sastels@users.noreply.github.com>
Date: Tue, 23 Jan 2024 14:39:28 -0500
Subject: [PATCH] Script to add rows to notification_history table (#2068)

---
 scripts/enlarge_db/README.md     | 22 +++++++++++++
 scripts/enlarge_db/enlarge_db.py | 55 ++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 scripts/enlarge_db/README.md
 create mode 100644 scripts/enlarge_db/enlarge_db.py

diff --git a/scripts/enlarge_db/README.md b/scripts/enlarge_db/README.md
new file mode 100644
index 0000000000..b36090390b
--- /dev/null
+++ b/scripts/enlarge_db/README.md
@@ -0,0 +1,22 @@
+# Enlarge DB
+
+## Purpose
+
+The purpose of this script is add rows to the notification_history table. This is useful in estimating how long database-related infrastructure operations will take when performed on a database the same size as that in production.
+
+## How to use
+
+The script should be run in the same environment as api. Locally this can be in the api repo devcontainer, while in AWS the api kubernetes pod would be preferred.
+
+To add 2000 rows to the table with a client_reference of "test2000" run
+
+```
+cd scripts/enlarge_db
+python enlarge_db.py -n 2000 -r test2000
+```
+
+The new notifications are added in batches to improve performance, with a default batch size of 10000. You may use a different batch with the `-c` parameter, for example
+
+```
+python enlarge_db.py -n 2000 -c 101 -r test2000x101
+```
diff --git a/scripts/enlarge_db/enlarge_db.py b/scripts/enlarge_db/enlarge_db.py
new file mode 100644
index 0000000000..649e6f3032
--- /dev/null
+++ b/scripts/enlarge_db/enlarge_db.py
@@ -0,0 +1,55 @@
+
+import argparse
+import sys
+from datetime import datetime
+from typing import List
+
+from flask import Flask
+
+sys.path.append("../..")
+from app import create_app, create_uuid, db  # noqa: E402
+from app.config import Config  # noqa: E402
+from app.models import NotificationHistory  # noqa: E402
+
+DEFAULT_CHUNK_SIZE = 10000
+
+
+def create_notifications(n: int, ref: str) -> List[NotificationHistory]:
+    notifications = [
+        NotificationHistory(
+            id=create_uuid(),
+            created_at=datetime.utcnow(),
+            template_id=Config.NEW_USER_EMAIL_VERIFICATION_TEMPLATE_ID,
+            template_version=1,
+            service_id=Config.NOTIFY_SERVICE_ID,
+            notification_type="email",
+            key_type='normal',
+            client_reference=ref,
+        )
+        for _ in range(n)
+    ]
+    return notifications
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--notifications", default=1, type=int, help="number of notifications to add to the notification_history table (default 1)")
+    parser.add_argument("-r", "--reference", default="manually created", type=str, help="client reference to use for the notifications (default 'manually created')")
+    parser.add_argument("-c", "--chunksize", default=DEFAULT_CHUNK_SIZE, type=int, help=f"chunk size for bulk_save_objects (default {DEFAULT_CHUNK_SIZE})")
+    args = parser.parse_args()
+
+    app = Flask("enlarge_db")
+    create_app(app)
+
+    for notifications_done in range(0, args.notifications, args.chunksize):
+        notifications = create_notifications(min(args.chunksize, args.notifications - notifications_done), args.reference)
+        print(f"Adding {len(notifications)} notifications to notification_history")
+        with app.app_context():
+            try:
+                db.session.bulk_save_objects(notifications)
+                db.session.commit()
+            except Exception as e:
+                print(f"Error adding notifications: {e}")
+                db.session.rollback()
+                sys.exit(1)
+        print(f"Done {notifications_done+len(notifications)} / {args.notifications}")