Sync on partner user creation + several fixes (#2214)

* Do not close session since it leads to orphan user object

* Redirect instead of render to avoid having to have a mailbox object

* On inital partner link/login trigger sync

* Update github action upload/artifact to v4

* Remove sys.exit used to test script locally

* Simplified script to update alias flags and note

.github/workflows/main.yml

@@ -109,7 +109,7 @@ jobs:
           GITHUB_ACTIONS_TEST: true
 
       - name: Archive code coverage results
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: code-coverage-report
           path: htmlcov

app/dashboard/views/mailbox.py

@@ -128,7 +128,7 @@ def mailbox_verify():
     except mailbox_utils.MailboxError as e:
         LOG.i(f"Cannot verify mailbox {mailbox_id} because of {e}")
         flash(f"Cannot verify mailbox: {e.msg}", "error")
-        return render_template("dashboard/mailbox_validation.html", mailbox=mailbox)
+        return redirect(url_for("dashboard.mailbox_route"))
     LOG.d("Mailbox %s is verified", mailbox)
     return render_template("dashboard/mailbox_validation.html", mailbox=mailbox)
 

app/proton/proton_callback_handler.py

@@ -2,9 +2,11 @@ from dataclasses import dataclass
 from enum import Enum
 from flask import url_for
 from typing import Optional
+import arrow
 
+from app import config
 from app.errors import LinkException
-from app.models import User, Partner
+from app.models import User, Partner, Job
 from app.proton.proton_client import ProtonClient, ProtonUser
 from app.account_linking import (
     process_login_case,
@@ -41,12 +43,21 @@ class ProtonCallbackHandler:
     def __init__(self, proton_client: ProtonClient):
         self.proton_client = proton_client
 
+    def _initial_alias_sync(self, user: User):
+        Job.create(
+            name=config.JOB_SEND_ALIAS_CREATION_EVENTS,
+            payload={"user_id": user.id},
+            run_at=arrow.now(),
+            commit=True,
+        )
+
     def handle_login(self, partner: Partner) -> ProtonCallbackResult:
         try:
             user = self.__get_partner_user()
             if user is None:
                 return generate_account_not_allowed_to_log_in()
             res = process_login_case(user, partner)
+            self._initial_alias_sync(res.user)
             return ProtonCallbackResult(
                 redirect_to_login=False,
                 flash_message=None,
@@ -75,6 +86,7 @@ class ProtonCallbackHandler:
             if user is None:
                 return generate_account_not_allowed_to_log_in()
             res = process_link_case(user, current_user, partner)
+            self._initial_alias_sync(res.user)
             return ProtonCallbackResult(
                 redirect_to_login=False,
                 flash_message="Account successfully linked",

email_handler.py

@@ -262,8 +262,6 @@ def get_or_create_contact(from_header: str, mail_from: str, alias: Alias) -> Con
 
             Session.commit()
         except IntegrityError:
-            # If the tx has been rolled back, the connection is borked. Force close to try to get a new one and start fresh
-            Session.close()
             LOG.info(
                 f"Contact with email {contact_email} for alias_id {alias_id} already existed, fetching from DB"
             )
@@ -818,7 +816,7 @@ def forward_email_to_mailbox(
 
     email_log = EmailLog.create(
         contact_id=contact.id,
-        user_id=user.id,
+        user_id=contact.user_id,
         mailbox_id=mailbox.id,
         alias_id=contact.alias_id,
         message_id=str(msg[headers.MESSAGE_ID]),

oneshot/alias_partner_set_flag_and_clear_note.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+import argparse
+import time
+
+from sqlalchemy import func
+from app.models import Alias
+from app.db import Session
+
+parser = argparse.ArgumentParser(
+    prog="Backfill alias", description="Update alias notes and backfill 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"Checking alias {alias_id_start} to {max_alias_id}")
+step = 1000
+noteSql = "(note = 'Created through Proton' or note = 'Created through partner Proton')"
+alias_query = f"UPDATE alias set note = NULL, flags = flags | :flag where id>=:start AND id<:end and {noteSql}"
+updated = 0
+start_time = time.time()
+for batch_start in range(alias_id_start, max_alias_id, step):
+    rows_done = Session.execute(
+        alias_query,
+        {
+            "start": batch_start,
+            "end": batch_start + step,
+            "flag": Alias.FLAG_PARTNER_CREATED,
+        },
+    )
+    updated += rows_done.rowcount
+    Session.commit()
+    elapsed = time.time() - start_time
+    time_per_alias = elapsed / (updated + 1)
+    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
+    print(
+        f"\rAlias {batch_start}/{max_alias_id} {updated} {hours_remaining:.2f}hrs remaining"
+    )
+print("")

tests/proton/test_proton_callback_handler.py

@@ -1,4 +1,6 @@
 from arrow import Arrow
+
+from app import config
 from app.account_linking import (
     SLPlan,
     SLPlanType,
@@ -8,7 +10,7 @@ from app.proton.proton_callback_handler import (
     ProtonCallbackHandler,
     generate_account_not_allowed_to_log_in,
 )
-from app.models import User, PartnerUser
+from app.models import User, PartnerUser, Job, JobState
 from app.proton.utils import get_proton_partner
 from app.utils import random_string
 from typing import Optional
@@ -23,6 +25,17 @@ class MockProtonClient(ProtonClient):
         return self.user
 
 
+def check_initial_sync_job(user: User):
+    for job in Job.yield_per_query(10).filter_by(
+        name=config.JOB_SEND_ALIAS_CREATION_EVENTS,
+        state=JobState.ready.value,
+    ):
+        if job.payload.get("user_id") == user.id:
+            Job.delete(job.id)
+            return
+    assert False
+
+
 def test_proton_callback_handler_unexistant_sl_user():
     email = random_email()
     name = random_string()
@@ -56,6 +69,7 @@ def test_proton_callback_handler_unexistant_sl_user():
     )
     assert partner_user is not None
     assert partner_user.external_user_id == external_id
+    check_initial_sync_job(res.user)
 
 
 def test_proton_callback_handler_existant_sl_user():
@@ -84,6 +98,7 @@ def test_proton_callback_handler_existant_sl_user():
     sa = PartnerUser.get_by(user_id=sl_user.id, partner_id=get_proton_partner().id)
     assert sa is not None
     assert sa.partner_email == user.email
+    check_initial_sync_job(res.user)
 
 
 def test_proton_callback_handler_none_user_login():