From 74b811dd35ce148848ddade7fb6c06201592a416 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?=
 <acasajus@users.noreply.github.com>
Date: Thu, 14 Mar 2024 11:11:50 +0100
Subject: [PATCH] Update oneshot commands (#2060)

* Update oneshot commands

* fix

* Fix test_load command

* Rename to avoid test executing it
---
 oneshot/emulate_dummy_load.py                 | 37 ++++++++++++
 oneshot/mark_aliases_as_partner_created.py    | 57 ++++++++++++++-----
 ...acts.py => replace_noreply_in_contacts.py} |  0
 3 files changed, 79 insertions(+), 15 deletions(-)
 create mode 100644 oneshot/emulate_dummy_load.py
 rename oneshot/{replace_noreply_in_cotnacts.py => replace_noreply_in_contacts.py} (100%)

diff --git a/oneshot/emulate_dummy_load.py b/oneshot/emulate_dummy_load.py
new file mode 100644
index 00000000..9766abf4
--- /dev/null
+++ b/oneshot/emulate_dummy_load.py
@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+import argparse
+import random
+import time
+
+from sqlalchemy import func
+
+from app import config
+from app.models import Alias, Contact
+from app.db import Session
+
+parser = argparse.ArgumentParser(
+    prog=f"Replace {config.NOREPLY}",
+    description=f"Replace {config.NOREPLY} from contacts reply email",
+)
+args = parser.parse_args()
+
+max_alias_id: int = Session.query(func.max(Alias.id)).scalar()
+
+start = time.time()
+tests = 1000
+for i in range(tests):
+    alias = (
+        Alias.filter(Alias.id > int(random.random() * max_alias_id))
+        .order_by(Alias.id.asc())
+        .limit(1)
+        .first()
+    )
+    contact = Contact.filter_by(alias_id=alias.id).order_by(Contact.id.asc()).first()
+    mailboxes = alias.mailboxes
+    user = alias.user
+    if i % 10:
+        print("{i} -> {alias.id}")
+
+end = time.time()
+time_taken = end - start
+print(f"Took {time_taken} -> {time_taken/tests} per test")
diff --git a/oneshot/mark_aliases_as_partner_created.py b/oneshot/mark_aliases_as_partner_created.py
index cf99db91..9c38fc00 100644
--- a/oneshot/mark_aliases_as_partner_created.py
+++ b/oneshot/mark_aliases_as_partner_created.py
@@ -1,29 +1,56 @@
 #!/usr/bin/env python3
 import argparse
+import time
 
+from sqlalchemy import func
 
-from app.log import LOG
 from app.models import Alias, SLDomain
 from app.db import Session
 
 parser = argparse.ArgumentParser(
     prog="Mark partner created aliases with the PARTNER_CREATED flag",
 )
+parser.add_argument(
+    "-s", "--start_alias_id", default=0, type=int, help="Initial alias_id"
+)
+parser.add_argument("-e", "--end_alias_id", default=0, type=int, help="Last alias_id")
+
 args = parser.parse_args()
+alias_id_start = args.start_alias_id
+max_alias_id = args.end_alias_id
+if max_alias_id == 0:
+    max_alias_id = Session.query(func.max(Alias.id)).scalar()
+
+print(f"Updating aliases from {alias_id_start} to {max_alias_id}")
 
 domains = SLDomain.filter(SLDomain.partner_id.isnot(None)).all()
+cond = [f"email like '%{domain.domain}'" for domain in domains]
+sql_or_cond = " OR ".join(cond)
+sql = f"UPDATE alias set flags = (flags | :flag) WHERE id >= :start and id<:end and flags & :flag = 0 and ({sql_or_cond})"
+print(sql)
 
-for domain in domains:
-    LOG.i(f"Checking aliases for domain {domain.domain}")
-    for alias in (
-        Alias.filter(
-            Alias.email.like(f"%{domain.domain}"),
-            Alias.flags.op("&")(Alias.FLAG_PARTNER_CREATED) == 0,
-        )
-        .enable_eagerloads(False)
-        .yield_per(100)
-        .all()
-    ):
-        alias.flags = alias.flags | Alias.FLAG_PARTNER_CREATED
-        LOG.i(f" * Updating {alias.email} to {alias.flags}")
-        Session.commit()
+step = 1000
+updated = 0
+start_time = time.time()
+for batch_start in range(alias_id_start, max_alias_id, step):
+    updated += Session.execute(
+        sql,
+        {
+            "start": batch_start,
+            "end": batch_start + step,
+            "flag": Alias.FLAG_PARTNER_CREATED,
+        },
+    ).rowcount
+    elapsed = time.time() - start_time
+    time_per_alias = elapsed / (batch_start - alias_id_start + step)
+    last_batch_id = batch_start + step
+    remaining = max_alias_id - last_batch_id
+    time_remaining = (max_alias_id - last_batch_id) * time_per_alias
+    hours_remaining = time_remaining / 3600.0
+    percent = int(
+        ((batch_start - alias_id_start) * 100) / (max_alias_id - alias_id_start)
+    )
+    print(
+        f"\rAlias {batch_start}/{max_alias_id} {percent}% {updated} updated {hours_remaining:.2f}hrs remaining"
+    )
+print(f"Updated aliases up to {max_alias_id}")
diff --git a/oneshot/replace_noreply_in_cotnacts.py b/oneshot/replace_noreply_in_contacts.py
similarity index 100%
rename from oneshot/replace_noreply_in_cotnacts.py
rename to oneshot/replace_noreply_in_contacts.py