joinedload folder and labels

inbox/mailsync/backends/gmail.py

@@ -449,7 +449,7 @@ def __deduplicate_message_object_creation(self, db_session, raw_messages, accoun
                 uid.update_flags(raw_message.flags)
                 uid.update_labels(raw_message.g_labels)
                 common.update_message_metadata(
-                    db_session, account, message_obj, uid.is_draft
+                    db_session, account.category_type, message_obj, uid.is_draft
                 )
                 db_session.commit()
 

inbox/mailsync/backends/imap/common.py

@@ -12,19 +12,21 @@
 
 """
 
+from collections.abc import Iterable
 from datetime import datetime
-from typing import List, Set
+from typing import List, Set, assert_never
 
 from sqlalchemy import bindparam, desc
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, joinedload
 from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.sql.expression import func
 
 from inbox.contacts.processing import update_contacts_from_message
 from inbox.crispin import RawMessage
 from inbox.logging import get_logger
 from inbox.models import Account, ActionLog, Folder, Message, MessageCategory
-from inbox.models.backends.imap import ImapFolderInfo, ImapUid
+from inbox.models.account import CategoryType
+from inbox.models.backends.imap import ImapFolderInfo, ImapUid, LabelItem
 from inbox.models.category import Category
 from inbox.models.session import session_scope
 from inbox.models.util import reconcile_message
@@ -74,7 +76,7 @@ def lastseenuid(account_id, session, folder_id):
 
 
 def update_message_metadata(
-    session: Session, account: Account, message: Message, is_draft: bool
+    session: Session, category_type: CategoryType, message: Message, is_draft: bool
 ) -> None:
     """Update the message's metadata"""
     # Sort imapuids in a way that the ones that were added later come last
@@ -88,24 +90,27 @@ def update_message_metadata(
     message.is_draft = is_draft
 
     sorted_categories: List[Category] = [
-        category for imapuid in sorted_imapuids for category in imapuid.categories
+        category
+        for imapuid in sorted_imapuids
+        for category in imapuid.get_categories(category_type)
     ]
 
     categories: Set[Category]
-    if account.category_type == "folder":
-        # For generic IMAP we want to deterministically select the last category.
-        # A message will always be in a single folder but it seems that for some
-        # on-prem servers we are not able to reliably detect when a message is moved
-        # between folders and we end up with many folders in our MySQL.
-        # Such message used to undeterministically appear in one of those folders
-        # (and in turn one category) depending on the order they were returned
-        # from the database. This makes it deterministic and more-correct because a message
-        # is likely in a folder (and category) it was added to last.
-        categories = {sorted_categories[-1]} if sorted_categories else set()
-    elif account.category_type == "label":
-        categories = set(sorted_categories)
-    else:
-        raise AssertionError("Unreachable")
+    match category_type:
+        case "folder":
+            # For generic IMAP we want to deterministically select the last category.
+            # A message will always be in a single folder but it seems that for some
+            # on-prem servers we are not able to reliably detect when a message is moved
+            # between folders and we end up with many folders in our MySQL.
+            # Such message used to undeterministically appear in one of those folders
+            # (and in turn one category) depending on the order they were returned
+            # from the database. This makes it deterministic and more-correct because a message
+            # is likely in a folder (and category) it was added to last.
+            categories = {sorted_categories[-1]} if sorted_categories else set()
+        case "label":
+            categories = set(sorted_categories)
+        case _:
+            assert_never(category_type)
 
     # Use a consistent time across creating categories, message updated_at
     # and the subsequent transaction that may be created.
@@ -193,12 +198,14 @@ def update_metadata(account_id, folder_id, folder_role, new_flags, session):
         if changed:
             change_count += 1
             is_draft = item.is_draft and folder_role in ["drafts", "all"]
-            update_message_metadata(session, account, item.message, is_draft)
+            update_message_metadata(
+                session, account.category_type, item.message, is_draft
+            )
             session.commit()
     log.info("Updated UID metadata", changed=change_count, out_of=len(new_flags))
 
 
-def remove_deleted_uids(account_id, folder_id, uids):
+def remove_deleted_uids(account_id: int, folder_id: int, uids: Iterable[int]) -> None:
     """
     Make sure you're holding a db write lock on the account. (We don't try
     to grab the lock in here in case the caller needs to put higher-level
@@ -207,7 +214,24 @@ def remove_deleted_uids(account_id, folder_id, uids):
     """
     if not uids:
         return
+
     deleted_uid_count = 0
+
+    with session_scope(account_id) as db_session:
+        category_type = Account.get(account_id, db_session).category_type
+
+    options = joinedload(ImapUid.message).joinedload(Message.imapuids)
+    match category_type:
+        case "folder":
+            options = options.joinedload(ImapUid.folder)
+        case "label":
+            options = options.options(
+                joinedload(ImapUid.folder),
+                joinedload("labelitems").joinedload(LabelItem.label),
+            )
+        case _:
+            assert_never(category_type)
+
     for uid in uids:
         # We do this one-uid-at-a-time because issuing many deletes within a
         # single database transaction is problematic. But loading many
@@ -228,7 +252,8 @@ def remove_deleted_uids(account_id, folder_id, uids):
                     ImapUid,
                     "FORCE INDEX (ix_imapuid_account_id_folder_id_msg_uid_desc)",
                 )
-                .first()
+                .options(options)
+                .one_or_none()
             )
             if imapuid is None:
                 continue
@@ -238,6 +263,8 @@ def remove_deleted_uids(account_id, folder_id, uids):
             db_session.delete(imapuid)
 
             if message is not None:
+                message.imapuids.remove(imapuid)
+
                 if not message.imapuids and message.is_draft:
                     # Synchronously delete drafts.
                     thread = message.thread
@@ -253,9 +280,8 @@ def remove_deleted_uids(account_id, folder_id, uids):
                     if thread is not None and not thread.messages:
                         db_session.delete(thread)
                 else:
-                    account = Account.get(account_id, db_session)
                     update_message_metadata(
-                        db_session, account, message, message.is_draft
+                        db_session, category_type, message, message.is_draft
                     )
                     if not message.imapuids:
                         # But don't outright delete messages. Just mark them as
@@ -324,7 +350,9 @@ def create_imap_message(
         is_draft = imapuid.is_draft and (
             folder.canonical_name == "drafts" or folder.canonical_name == "all"
         )
-        update_message_metadata(db_session, account, new_message, is_draft)
+        update_message_metadata(
+            db_session, account.category_type, new_message, is_draft
+        )
 
     update_contacts_from_message(db_session, new_message, account.namespace.id)
 

inbox/models/backends/imap.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import List, Set
+from typing import List, assert_never
 
 from sqlalchemy import (
     BigInteger,
@@ -20,7 +20,7 @@
 from sqlalchemy.sql.expression import false
 
 from inbox.logging import get_logger
-from inbox.models.account import Account
+from inbox.models.account import Account, CategoryType
 from inbox.models.base import MailSyncBase
 from inbox.models.category import Category
 from inbox.models.folder import Folder
@@ -209,11 +209,16 @@ def update_labels(self, new_labels: List[str]) -> None:
     def namespace(self):
         return self.imapaccount.namespace
 
-    @property
-    def categories(self) -> Set[Category]:
-        categories = {label.category for label in self.labels}
-        categories.add(self.folder.category)
-        return categories
+    def get_categories(self, category_type: CategoryType) -> set[Category]:
+        match category_type:
+            case "label":
+                return {label.category for label in self.labels} | {
+                    self.folder.category
+                }
+            case "folder":
+                return {self.folder.category}
+            case _:
+                assert_never(category_type)
 
     __table_args__ = (
         UniqueConstraint("folder_id", "msg_uid", "account_id"),

tests/imap/test_labels.py

@@ -32,7 +32,9 @@ def folder_and_message_maps(db, default_account):
         # Create a message in the folder
         message = add_fake_message(db.session, default_account.namespace.id, thread)
         add_fake_imapuid(db.session, default_account.id, message, folder, 13)
-        update_message_metadata(db.session, default_account, message, False)
+        update_message_metadata(
+            db.session, default_account.category_type, message, False
+        )
         db.session.commit()
         folder_map[name] = folder
         message_map[name] = message
@@ -46,7 +48,7 @@ def add_inbox_label(db, default_account, message):
     imapuid.update_labels(["\\Inbox"])
     db.session.commit()
     assert {c.name for c in imapuid.categories} == {"all", "inbox"}
-    update_message_metadata(db.session, default_account, message, False)
+    update_message_metadata(db.session, default_account.category_type, message, False)
     db.session.commit()
     return message
 
@@ -58,7 +60,7 @@ def add_custom_label(db, default_account, message):
     imapuid.update_labels(["<3"])
     db.session.commit()
     assert {c.name for c in imapuid.categories} == {existing, ""}
-    update_message_metadata(db.session, default_account, message, False)
+    update_message_metadata(db.session, default_account.category_type, message, False)
     db.session.commit()
     return message
 

tests/imap/test_update_metadata.py

@@ -68,7 +68,7 @@ def test_update_categories_when_actionlog_entry_missing(
 ):
     message.categories_changes = True
     db.session.commit()
-    update_message_metadata(db.session, imapuid.account, message, False)
+    update_message_metadata(db.session, imapuid.account.category_type, message, False)
     assert message.categories == {imapuid.folder.category}
 
 
@@ -112,7 +112,7 @@ def test_categories_from_multiple_imap_folders(
         imapuid.updated_at = imapuid.updated_at + datetime.timedelta(seconds=delay)
         db.session.commit()
 
-    update_message_metadata(db.session, generic_account, message, False)
+    update_message_metadata(db.session, generic_account.category_type, message, False)
     assert {category.name for category in message.categories} == categories
 
     delete_imapuids(db.session)