Skip to content

Commit

Permalink
Merge branch 'main' into task/read-from-last-used
Browse files Browse the repository at this point in the history
  • Loading branch information
jzbahrai authored Jan 23, 2024
2 parents df71f36 + 0f36abb commit a11752c
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 0 deletions.
22 changes: 22 additions & 0 deletions scripts/enlarge_db/README.md
Original file line number Diff line number Diff line change
@@ -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
```
55 changes: 55 additions & 0 deletions scripts/enlarge_db/enlarge_db.py
Original file line number Diff line number Diff line change
@@ -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}")

0 comments on commit a11752c

Please sign in to comment.