Catch missed txn rollbacks (#1094)

* catch missed rollbacks

* change approach

* revert back to simple way

* test

* add to bsky package as well

* improve handler & tests

packages/bsky/src/db/index.ts

@@ -62,7 +62,14 @@ export class Database {
       .transaction()
       .execute(async (txn) => {
         const dbTxn = new Database(txn, this.cfg)
-        const txRes = await fn(dbTxn).finally(() => leakyTxPlugin.endTx())
+        const txRes = await fn(dbTxn)
+          .catch(async (err) => {
+            leakyTxPlugin.endTx()
+            // ensure that all in-flight queries are flushed & the connection is open
+            await dbTxn.db.getExecutor().provideConnection(async () => {})
+            throw err
+          })
+          .finally(() => leakyTxPlugin.endTx())
         return txRes
       })
     return res

packages/bsky/tests/db.test.ts

@@ -125,6 +125,43 @@ describe('db', () => {
 
       expect(res.length).toBe(0)
     })
+
+    it('ensures all inflight querys are rolled back', async () => {
+      let promise: Promise<unknown> | undefined = undefined
+      const names: string[] = []
+      try {
+        await db.transaction(async (dbTxn) => {
+          const queries: Promise<unknown>[] = []
+          for (let i = 0; i < 20; i++) {
+            const name = `user${i}`
+            const query = dbTxn.db
+              .insertInto('actor')
+              .values({
+                handle: name,
+                did: name,
+                indexedAt: 'bad-date',
+              })
+              .execute()
+            names.push(name)
+            queries.push(query)
+          }
+          promise = Promise.allSettled(queries)
+          throw new Error()
+        })
+      } catch (err) {
+        expect(err).toBeDefined()
+      }
+      if (promise) {
+        await promise
+      }
+
+      const res = await db.db
+        .selectFrom('actor')
+        .selectAll()
+        .where('did', 'in', names)
+        .execute()
+      expect(res.length).toBe(0)
+    })
   })
 
   describe('Leader', () => {

packages/pds/src/db/index.ts

@@ -179,7 +179,14 @@ export class Database {
       .transaction()
       .execute(async (txn) => {
         const dbTxn = new Database(txn, this.cfg, this.channels)
-        const txRes = await fn(dbTxn).finally(() => leakyTxPlugin.endTx())
+        const txRes = await fn(dbTxn)
+          .catch(async (err) => {
+            leakyTxPlugin.endTx()
+            // ensure that all in-flight queries are flushed & the connection is open
+            await dbTxn.db.getExecutor().provideConnection(async () => {})
+            throw err
+          })
+          .finally(() => leakyTxPlugin.endTx())
         txMsgs = dbTxn.txChannelMsgs
         return { txRes, dbTxn }
       })

packages/pds/tests/db.test.ts

@@ -129,6 +129,43 @@ describe('db', () => {
 
       expect(res.length).toBe(0)
     })
+
+    it('ensures all inflight querys are rolled back', async () => {
+      let promise: Promise<unknown> | undefined = undefined
+      const names: string[] = []
+      try {
+        await db.transaction(async (dbTxn) => {
+          const queries: Promise<unknown>[] = []
+          for (let i = 0; i < 20; i++) {
+            const name = `user${i}`
+            const query = dbTxn.db
+              .insertInto('repo_root')
+              .values({
+                root: name,
+                did: name,
+                indexedAt: 'bad-date',
+              })
+              .execute()
+            names.push(name)
+            queries.push(query)
+          }
+          promise = Promise.allSettled(queries)
+          throw new Error()
+        })
+      } catch (err) {
+        expect(err).toBeDefined()
+      }
+      if (promise) {
+        await promise
+      }
+
+      const res = await db.db
+        .selectFrom('repo_root')
+        .selectAll()
+        .where('did', 'in', names)
+        .execute()
+      expect(res.length).toBe(0)
+    })
   })
 
   describe('Leader', () => {