add nested transaction with savepoints

src/db/connection.cr

@@ -84,5 +84,20 @@ module DB
     def perform_rollback_transaction
       self.unprepared.exec "ROLLBACK"
     end
+
+    # :nodoc:
+    def perform_create_savepoint(name)
+      self.unprepared.exec "SAVEPOINT #{name}"
+    end
+
+    # :nodoc:
+    def perform_release_savepoint(name)
+      self.unprepared.exec "RELEASE SAVEPOINT #{name}"
+    end
+
+    # :nodoc:
+    def perform_rollback_savepoint(name)
+      self.unprepared.exec "ROLLBACK TO #{name}"
+    end
   end
 end

src/db/transaction.cr

@@ -1,6 +1,7 @@
 module DB
   abstract class Transaction
     include Disposable
+    include BeginTransaction
 
     abstract def connection : Connection
 
@@ -16,28 +17,102 @@ module DB
       raise DB::Error.new("Transaction already closed") if closed?
       close
     end
+
+    abstract def release_from_nested_transaction
   end
 
   class TopLevelTransaction < Transaction
-    # :nodoc:
     getter connection
+    # :nodoc:
+    property savepoint_name : String? = nil
 
     def initialize(@connection : Connection)
+      @nested_transaction = false
       @connection.perform_begin_transaction
     end
 
     def commit
       @connection.perform_commit_transaction
-      close!
+      super
     end
 
     def rollback
       @connection.perform_rollback_transaction
-      close!
+      super
     end
 
     protected def do_close
       connection.release_from_transaction
     end
+
+    def begin_transaction : Transaction
+      raise DB::Error.new("There is an existing nested transaction in this transaction") if @nested_transaction
+      @nested_transaction = true
+      create_save_point_transaction(self)
+    end
+
+    # :nodoc:
+    def create_save_point_transaction(parent : Transaction) : SavePointTransaction
+      # TODO should we wrap this in a mutex?
+      previous_savepoint = @savepoint_name
+      savepoint_name = if previous_savepoint
+                         previous_savepoint.succ
+                       else
+                         # random prefix to avoid determinism
+                         "crystal_#{Random.rand(10_000)}_00001"
+                       end
+
+      @savepoint_name = savepoint_name
+
+      create_save_point_transaction(parent, savepoint_name)
+    end
+
+    protected def create_save_point_transaction(parent : Transaction, savepoint_name : String) : SavePointTransaction
+      SavePointTransaction.new(parent, savepoint_name)
+    end
+
+    # :nodoc:
+    def release_from_nested_transaction
+      @nested_transaction = false
+    end
+  end
+
+  class SavePointTransaction < Transaction
+    getter connection : Connection
+
+    def initialize(@parent : Transaction,
+                   @savepoint_name : String)
+      @nested_transaction = false
+      @connection = @parent.connection
+      @connection.perform_create_savepoint(@savepoint_name)
+    end
+
+    def commit
+      @connection.perform_release_savepoint(@savepoint_name)
+      super
+    end
+
+    def rollback
+      @connection.perform_rollback_savepoint(@savepoint_name)
+      super
+    end
+
+    protected def do_close
+      @parent.release_from_nested_transaction
+    end
+
+    def begin_transaction : Transaction
+      raise DB::Error.new("There is an existing nested transaction in this transaction") if @nested_transaction
+      @nested_transaction = true
+      create_save_point_transaction(self)
+    end
+
+    def create_save_point_transaction(parent : Transaction)
+      @parent.create_save_point_transaction(parent)
+    end
+
+    def release_from_nested_transaction
+      @nested_transaction = false
+    end
   end
 end