From 9db36a6ab163e0072b38e9dba7b444afe30d03f1 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 8 Jan 2025 16:06:16 +0000
Subject: [PATCH] Add helper functions

---
 synapse/storage/databases/main/state.py | 32 +++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/synapse/storage/databases/main/state.py b/synapse/storage/databases/main/state.py
index 101a82c509..d785cea9a3 100644
--- a/synapse/storage/databases/main/state.py
+++ b/synapse/storage/databases/main/state.py
@@ -149,6 +149,38 @@ class StateGroupWorkerStore(EventsWorkerStore, SQLBaseStore):
             "_advance_state_epoch", advance_state_epoch_txn, db_autocommit=True
         )
 
+    @cached()
+    async def is_state_group_pending_deletion_before(
+        self, state_epoch: int, state_group: int
+    ) -> bool:
+        """Check if a state group is marked as pending deletion in a previous
+        epoch, but does not check the current epoch."""
+
+        def is_state_group_pending_deletion_before_txn(txn: LoggingTransaction) -> bool:
+            sql = """
+                SELECT 1 FROM state_groups_pending_deletion
+                WHERE state_epoch < ? AND state_group = ?
+            """
+            txn.execute(sql, (state_epoch, state_group))
+
+            return txn.fetchone() is not None
+
+        return await self.db_pool.runInteraction(
+            "is_state_group_pending_deletion_before",
+            is_state_group_pending_deletion_before_txn,
+        )
+
+    async def mark_state_group_as_used(self, state_group: int) -> None:
+        """Mark that a given state group is used"""
+
+        # TODO: Also assert that the state group hasn't advanced too much
+
+        await self.db_pool.simple_delete(
+            table="state_groups_pending_deletion",
+            keyvalues={"state_group": state_group},
+            desc="mark_state_group_as_used",
+        )
+
     def process_replication_rows(
         self,
         stream_name: str,