From 703e6b501f2f2b91693b255ef346ed829d602568 Mon Sep 17 00:00:00 2001
From: Harminder Virk <virk.officials@gmail.com>
Date: Tue, 27 Feb 2024 15:38:44 +0530
Subject: [PATCH] feat: add model.lockForUpdate method to lock the model row
 for updates

---
 src/orm/base_model/index.ts |  24 +++++++
 src/types/model.ts          |  16 +++++
 test/orm/base_model.spec.ts | 124 ++++++++++++++++++++++++++++++++++++
 3 files changed, 164 insertions(+)

diff --git a/src/orm/base_model/index.ts b/src/orm/base_model/index.ts
index 1efbce2d..1ef649b2 100644
--- a/src/orm/base_model/index.ts
+++ b/src/orm/base_model/index.ts
@@ -1845,6 +1845,30 @@ class BaseModelImpl implements LucidRow {
     return this
   }
 
+  /**
+   * The lockForUpdate method re-fetches the model instance from
+   * the database and locks the row to perform an update. The
+   * provided callback receives a fresh user instance and should
+   * use that to perform an update.
+   */
+  async lockForUpdate<T>(callback: (user: this) => Promise<T> | T): Promise<T> {
+    const Model = this.constructor as LucidModel
+    const queryClient = Model.$adapter.modelClient(this)
+
+    return managedTransaction<T>(queryClient, async (trx) => {
+      const user = await Model.query({ client: trx })
+        .forUpdate()
+        .where(Model.primaryKey, this.$primaryKeyValue)
+        .first()
+
+      if (!user) {
+        throw new errors.E_MODEL_DELETED()
+      }
+
+      return callback(user as this)
+    })
+  }
+
   /**
    * Perform delete by issuing a delete request on the adapter
    */
diff --git a/src/types/model.ts b/src/types/model.ts
index a0bdc286..65efb4f6 100644
--- a/src/types/model.ts
+++ b/src/types/model.ts
@@ -619,7 +619,23 @@ export interface LucidRow {
    * Actions to perform on the instance
    */
   save(): Promise<this>
+
+  /**
+   * The lockForUpdate method re-fetches the model instance from
+   * the database and locks the row to perform an update. The
+   * provided callback receives a fresh user instance and should
+   * use that to perform an update.
+   */
+  lockForUpdate<T>(callback: (user: this) => Promise<T> | T): Promise<T>
+
+  /**
+   * Perform delete by issuing a delete request on the adapter
+   */
   delete(): Promise<void>
+
+  /**
+   * Reload/Refresh the model instance
+   */
   refresh(): Promise<this>
 
   /**
diff --git a/test/orm/base_model.spec.ts b/test/orm/base_model.spec.ts
index 205c4b1b..27f2f4bb 100644
--- a/test/orm/base_model.spec.ts
+++ b/test/orm/base_model.spec.ts
@@ -7926,3 +7926,127 @@ test.group('Base model | inheritance', (group) => {
     assert.isFalse(MyBaseModel.$hooks.has('before:create', hook2))
   })
 })
+
+test.group('Base Model | lockForUpdate', (group) => {
+  group.setup(async () => {
+    await setup()
+  })
+
+  group.teardown(async () => {
+    await cleanupTables()
+  })
+
+  test('lock model row for update', async ({ fs, assert }) => {
+    const app = new AppFactory().create(fs.baseUrl, () => {})
+    await app.init()
+    const db = getDb()
+    const adapter = ormAdapter(db)
+
+    const BaseModel = getBaseModel(adapter)
+
+    class User extends BaseModel {
+      @column({ isPrimary: true })
+      declare id: number
+
+      @column()
+      declare username: string
+
+      @column()
+      declare email: string
+
+      @column()
+      declare points: number
+    }
+
+    const user = await User.create({ email: 'foo@bar.com', username: 'virk', points: 0 })
+    await Promise.all([
+      user.lockForUpdate(async (freshUser) => {
+        freshUser.points = freshUser.points + 1
+        await freshUser.save()
+      }),
+      user.lockForUpdate(async (freshUser) => {
+        freshUser.points = freshUser.points + 1
+        await freshUser.save()
+      }),
+    ])
+
+    await user.refresh()
+    assert.equal(user.points, 2)
+  })
+
+  test('re-use initial model transaction', async ({ fs, assert }) => {
+    const app = new AppFactory().create(fs.baseUrl, () => {})
+    await app.init()
+    const db = getDb()
+    const adapter = ormAdapter(db)
+
+    const BaseModel = getBaseModel(adapter)
+
+    class User extends BaseModel {
+      @column({ isPrimary: true })
+      declare id: number
+
+      @column()
+      declare username: string
+
+      @column()
+      declare email: string
+
+      @column()
+      declare points: number
+    }
+
+    const user = await User.create({ email: 'foo@bar.com', username: 'virk', points: 0 })
+    const trx = await db.transaction()
+    user.useTransaction(trx)
+
+    await Promise.all([
+      user.lockForUpdate(async (freshUser) => {
+        freshUser.points = freshUser.points + 1
+        await freshUser.save()
+      }),
+      user.lockForUpdate(async (freshUser) => {
+        freshUser.points = freshUser.points + 1
+        await freshUser.save()
+      }),
+    ])
+
+    assert.isFalse(trx.isCompleted)
+    await trx.rollback()
+
+    await user.refresh()
+    assert.equal(user.points, 0)
+  })
+
+  test('throw error when row is missing', async ({ fs, assert }) => {
+    const app = new AppFactory().create(fs.baseUrl, () => {})
+    await app.init()
+    const db = getDb()
+    const adapter = ormAdapter(db)
+
+    const BaseModel = getBaseModel(adapter)
+
+    class User extends BaseModel {
+      @column({ isPrimary: true })
+      declare id: number
+
+      @column()
+      declare username: string
+
+      @column()
+      declare email: string
+
+      @column()
+      declare points: number
+    }
+
+    const user = await User.create({ email: 'foo@bar.com', username: 'virk', points: 0 })
+    await user.delete()
+    await assert.rejects(() =>
+      user.lockForUpdate(async (freshUser) => {
+        freshUser.points = freshUser.points + 1
+        await freshUser.save()
+      })
+    )
+  })
+})