Adding table locking. Fixes #594

src/avram/database.cr

@@ -84,6 +84,21 @@ abstract class Avram::Database
     end
   end
 
+  # Creates a lock on the table in `mode`
+  # ```
+  # AppDatabase.with_lock_on(User, mode: :row_exclusive) do
+  #   user = UserQuery.new.id(1).for_update.first
+  #   SaveUser.update!(user, name: "New Name")
+  # end
+  # ```
+  def self.with_lock_on(model : Avram::Model.class, mode : Avram::TableLockMode, &)
+    exec("BEGIN")
+    exec("LOCK TABLE #{model.table_name} IN #{mode} MODE")
+    yield
+  ensure
+    exec("END")
+  end
+
   # Methods without a block
   {% for crystal_db_alias in [:exec, :scalar, :query, :query_all, :query_one, :query_one?] %}
     # Same as crystal-db's `DB::QueryMethods#{{ crystal_db_alias.id }}` but with instrumentation

src/avram/query_builder.cr

@@ -14,6 +14,7 @@ class Avram::QueryBuilder
   @prepared_statement_placeholder = 0
   @distinct : Bool = false
   @delete : Bool = false
+  @for_update : Bool = false
 
   def initialize(@table)
   end
@@ -116,7 +117,7 @@ class Avram::QueryBuilder
   end
 
   private def sql_condition_clauses
-    [joins_sql, wheres_sql, group_sql, order_sql, limit_sql, offset_sql]
+    [joins_sql, wheres_sql, group_sql, order_sql, limit_sql, offset_sql, locking_sql]
   end
 
   def delete : self
@@ -158,6 +159,11 @@ class Avram::QueryBuilder
     self
   end
 
+  def for_update : self
+    @for_update = true
+    self
+  end
+
   def order_by(order : OrderByClause) : self
     reset_order if ordered_randomly?
     @orders << order
@@ -384,4 +390,10 @@ class Avram::QueryBuilder
   private def delete_sql : String
     "DELETE FROM #{table}"
   end
+
+  private def locking_sql : String?
+    if @for_update
+      "FOR UPDATE"
+    end
+  end
 end

src/avram/queryable.cr

@@ -196,6 +196,10 @@ module Avram::Queryable(T)
     clone.tap &.query.offset(amount)
   end
 
+  def for_update : self
+    clone.tap &.query.for_update
+  end
+
   def first? : T?
     with_ordered_query
       .limit(1)

src/avram/table_lock_mode.cr

@@ -0,0 +1,16 @@
+module Avram
+  enum TableLockMode
+    ACCESS_SHARE
+    ROW_SHARE
+    ROW_EXCLUSIVE
+    SHARE_UPDATE_EXCLUSIVE
+    SHARE
+    SHARE_ROW_EXCLUSIVE
+    EXCLUSIVE
+    ACCESS_EXCLUSIVE
+
+    def to_s
+      member_name.to_s.gsub("_", " ")
+    end
+  end
+end