diff --git a/backend/app/main.py b/backend/app/main.py index 1ab8841..23d7e98 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -30,6 +30,37 @@ settings = get_settings() inbound_server = InboundSMTPServer() +async def ensure_mongo_index(collection, keys, **options) -> None: + """Create an index, replacing any same-name or same-key index with different options.""" + key_list = [(field, direction) for field, direction in keys] + options = dict(options) + name = options.get("name") + if not name: + name = "_".join(f"{field}_{direction}" for field, direction in key_list) + options["name"] = name + + desired_unique = bool(options.get("unique")) + desired_sparse = bool(options.get("sparse")) + + async for index in collection.list_indexes(): + if index["name"] == "_id_": + continue + index_key = list(index["key"].items()) + same_key = index_key == key_list + same_name = index["name"] == name + if not same_key and not same_name: + continue + same_options = ( + bool(index.get("unique")) == desired_unique + and bool(index.get("sparse")) == desired_sparse + ) + if same_key and same_name and same_options: + return + await collection.drop_index(index["name"]) + + await collection.create_index(keys, **options) + + async def backfill_missing_thread_ids(database) -> None: cursor = database.mailbox_entries.find({"thread_id": None}) async for entry in cursor: @@ -141,24 +172,28 @@ async def lifespan(_: FastAPI): async with SessionFactory() as session: await ensure_platform_domain(session) await session.commit() - await mongo_db.emails.create_index([("recipients", 1), ("created_at", -1)]) - await mongo_db.emails.create_index([("cc", 1), ("created_at", -1)]) - await mongo_db.emails.create_index([("bcc", 1), ("created_at", -1)]) - await mongo_db.emails.create_index([("sender", 1), ("created_at", -1)]) - await mongo_db.emails.create_index([("account_id", 1), ("created_at", -1)]) - await mongo_db.messages.create_index( + await ensure_mongo_index(mongo_db.emails, [("recipients", 1), ("created_at", -1)]) + await ensure_mongo_index(mongo_db.emails, [("cc", 1), ("created_at", -1)]) + await ensure_mongo_index(mongo_db.emails, [("bcc", 1), ("created_at", -1)]) + await ensure_mongo_index(mongo_db.emails, [("sender", 1), ("created_at", -1)]) + await ensure_mongo_index(mongo_db.emails, [("account_id", 1), ("created_at", -1)]) + await ensure_mongo_index( + mongo_db.messages, [("delivery_key", 1)], unique=True, sparse=True, ) - await mongo_db.message_tombstones.create_index( + await ensure_mongo_index( + mongo_db.message_tombstones, [("delivery_key", 1)], unique=True, ) - await mongo_db.messages.create_index( - [("subject", "text"), ("snippet", "text"), ("sender.name", "text")] + await ensure_mongo_index( + mongo_db.messages, + [("subject", "text"), ("snippet", "text"), ("sender.name", "text")], ) - await mongo_db.mailbox_entries.create_index( + await ensure_mongo_index( + mongo_db.mailbox_entries, [ ("account_id", 1), ("mailbox", 1), @@ -166,22 +201,27 @@ async def lifespan(_: FastAPI): ("deletion_state", 1), ("created_at", -1), ("_id", -1), - ] + ], ) - await mongo_db.mailbox_entries.create_index( - [("account_id", 1), ("is_read", 1), ("created_at", -1), ("_id", -1)] + await ensure_mongo_index( + mongo_db.mailbox_entries, + [("account_id", 1), ("is_read", 1), ("created_at", -1), ("_id", -1)], ) - await mongo_db.mailbox_entries.create_index( + await ensure_mongo_index( + mongo_db.mailbox_entries, [("account_id", 1), ("message_id", 1), ("mailbox", 1)], unique=True, ) - await mongo_db.mailbox_entries.create_index( - [("account_id", 1), ("is_starred", 1), ("created_at", -1), ("_id", -1)] + await ensure_mongo_index( + mongo_db.mailbox_entries, + [("account_id", 1), ("is_starred", 1), ("created_at", -1), ("_id", -1)], ) - await mongo_db.sticky_notes.create_index( - [("account_id", 1), ("updated_at", -1)] + await ensure_mongo_index( + mongo_db.sticky_notes, + [("account_id", 1), ("updated_at", -1)], ) - await mongo_db.mailbox_entries.create_index( + await ensure_mongo_index( + mongo_db.mailbox_entries, [ ("account_id", 1), ("sender.name", "text"), @@ -191,37 +231,37 @@ async def lifespan(_: FastAPI): ], name="mailbox_search", ) - await mongo_db.messages.create_index( - [("message_id", 1)], - sparse=True, - ) - await mongo_db.messages.create_index( - [("message_id_folded", 1)], - sparse=True, - ) - await mongo_db.messages.create_index( + await ensure_mongo_index(mongo_db.messages, [("message_id", 1)], sparse=True) + await ensure_mongo_index(mongo_db.messages, [("message_id_folded", 1)], sparse=True) + await ensure_mongo_index( + mongo_db.messages, [("generated_from_handoff_message_id", 1)], sparse=True, unique=True, ) - await mongo_db.messages.create_index( + await ensure_mongo_index( + mongo_db.messages, [("mirrored_from_external_message_id", 1)], sparse=True, unique=True, ) - await mongo_db.messages.create_index( + await ensure_mongo_index( + mongo_db.messages, [("handoff_delivery_key", 1)], sparse=True, unique=True, ) - await mongo_db.mailbox_entries.create_index( - [("account_id", 1), ("thread_id", 1), ("event_at", -1), ("_id", -1)] + await ensure_mongo_index( + mongo_db.mailbox_entries, + [("account_id", 1), ("thread_id", 1), ("event_at", -1), ("_id", -1)], ) - await mongo_db.email_threads.create_index( - [("account_id", 1), ("last_message_at", -1), ("_id", -1)] + await ensure_mongo_index( + mongo_db.email_threads, + [("account_id", 1), ("last_message_at", -1), ("_id", -1)], ) - await mongo_db.human_handoff_rules.create_index( - [("account_id", 1), ("updated_at", -1)] + await ensure_mongo_index( + mongo_db.human_handoff_rules, + [("account_id", 1), ("updated_at", -1)], ) account_ids = await mongo_db.human_handoff_rules.distinct("account_id") for account_id in account_ids: @@ -236,45 +276,56 @@ async def lifespan(_: FastAPI): {"_id": rule["_id"]}, {"$set": {"priority": index}}, ) - await mongo_db.human_handoff_rules.create_index( + await ensure_mongo_index( + mongo_db.human_handoff_rules, [("account_id", 1), ("priority", 1)], unique=True, ) - await mongo_db.human_handoff_rules.create_index( + await ensure_mongo_index( + mongo_db.human_handoff_rules, [("account_id", 1), ("name_key", 1)], unique=True, sparse=True, ) - await mongo_db.human_handoffs.create_index( - [("account_id", 1), ("updated_at", -1)] + await ensure_mongo_index( + mongo_db.human_handoffs, + [("account_id", 1), ("updated_at", -1)], ) - await mongo_db.human_handoffs.create_index( + await ensure_mongo_index( + mongo_db.human_handoffs, [("account_id", 1), ("external_thread_id", 1)], unique=True, ) - await mongo_db.human_handoffs.create_index( + await ensure_mongo_index( + mongo_db.human_handoffs, [("account_id", 1), ("handoff_thread_id", 1)], unique=True, ) - await mongo_db.handoff_message_links.create_index( - [("handoff_id", 1), ("created_at", -1)] + await ensure_mongo_index( + mongo_db.handoff_message_links, + [("handoff_id", 1), ("created_at", -1)], ) - await mongo_db.handoff_message_links.create_index( + await ensure_mongo_index( + mongo_db.handoff_message_links, [("handoff_id", 1), ("source_message_id", 1), ("relationship", 1)], unique=True, sparse=True, ) - await mongo_db.handoff_delivery_jobs.create_index( + await ensure_mongo_index( + mongo_db.handoff_delivery_jobs, [("idempotency_key", 1)], unique=True, ) - await mongo_db.handoff_delivery_jobs.create_index( - [("status", 1), ("next_attempt_at", 1), ("created_at", 1)] + await ensure_mongo_index( + mongo_db.handoff_delivery_jobs, + [("status", 1), ("next_attempt_at", 1), ("created_at", 1)], ) - await mongo_db.mailbox_entries.create_index( - [("account_id", 1), ("message_scope", 1), ("mailbox", 1), ("event_at", -1)] + await ensure_mongo_index( + mongo_db.mailbox_entries, + [("account_id", 1), ("message_scope", 1), ("mailbox", 1), ("event_at", -1)], ) - await mongo_db.messages.create_index( + await ensure_mongo_index( + mongo_db.messages, [("message_id", 1), ("message_scope", 1)], sparse=True, )