add nested transaction with savepoints

spec/dummy_driver.cr

@@ -68,6 +68,29 @@ class DummyDriver < DB::Driver
       super
       @rolledback = true
     end
+
+    protected def create_save_point_transaction(parent, savepoint_name : String)
+      DummySavePointTransaction.new(parent, savepoint_name)
+    end
+  end
+
+  class DummySavePointTransaction < DB::SavePointTransaction
+    getter committed = false
+    getter rolledback = false
+
+    def initialize(parent, savepoint_name)
+      super(parent, savepoint_name)
+    end
+
+    def commit
+      super
+      @committed = true
+    end
+
+    def rollback
+      super
+      @rolledback = true
+    end
   end
 
   class DummyStatement < DB::Statement

spec/save_point_transaction_spec.cr

@@ -0,0 +1,127 @@
+require "./spec_helper"
+
+private class FooException < Exception
+end
+
+private def with_dummy_top_transaction
+  with_dummy_connection do |cnn|
+    cnn.transaction do |tx|
+      yield tx.as(DummyDriver::DummyTransaction), cnn
+    end
+  end
+end
+
+private def with_dummy_nested_transaction
+  with_dummy_connection do |cnn|
+    cnn.transaction do |tx|
+      tx.transaction do |nested|
+        yield nested.as(DummyDriver::DummySavePointTransaction), cnn
+      end
+    end
+  end
+end
+
+describe DB::SavePointTransaction do
+  {% for context in [:with_dummy_top_transaction, :with_dummy_nested_transaction] %}
+  describe "{{context.id}}" do
+    it "begin/commit transaction from parent transaction" do
+      {{context.id}} do |parent_tx|
+        tx = parent_tx.begin_transaction
+        tx.commit
+      end
+    end
+
+    it "begin/rollback transaction from parent transaction" do
+      {{context.id}} do |parent_tx|
+        tx = parent_tx.begin_transaction
+        tx.rollback
+      end
+    end
+
+    it "raise if begin over existing transaction" do
+      {{context.id}} do |parent_tx|
+        parent_tx.begin_transaction
+        expect_raises(DB::Error, "There is an existing nested transaction in this transaction") do
+          parent_tx.begin_transaction
+        end
+      end
+    end
+
+    it "allow sequential transactions" do
+      {{context.id}} do |parent_tx|
+        tx = parent_tx.begin_transaction
+        tx.rollback
+
+        tx = parent_tx.begin_transaction
+        tx.commit
+      end
+    end
+
+    it "transaction with block from parent transaction should be committed" do
+      t = uninitialized DummyDriver::DummySavePointTransaction
+
+      with_witness do |w|
+        {{context.id}} do |parent_tx|
+          parent_tx.transaction do |tx|
+            if tx.is_a?(DummyDriver::DummySavePointTransaction)
+              t = tx
+              w.check
+            end
+          end
+        end
+      end
+
+      t.committed.should be_true
+      t.rolledback.should be_false
+    end
+  end
+  {% end %}
+
+  it "only nested transaction with block from parent transaction should be rolledback if raise DB::Rollback" do
+    top = uninitialized DummyDriver::DummyTransaction
+    t = uninitialized DummyDriver::DummySavePointTransaction
+
+    with_witness do |w|
+      with_dummy_top_transaction do |parent_tx|
+        top = parent_tx
+        parent_tx.transaction do |tx|
+          if tx.is_a?(DummyDriver::DummySavePointTransaction)
+            t = tx
+            w.check
+          end
+          raise DB::Rollback.new
+        end
+      end
+    end
+
+    t.rolledback.should be_true
+    t.committed.should be_false
+
+    top.rolledback.should be_false
+    top.committed.should be_true
+  end
+
+  it "only nested transaction with block from parent nested transaction should be rolledback if raise DB::Rollback" do
+    top = uninitialized DummyDriver::DummySavePointTransaction
+    t = uninitialized DummyDriver::DummySavePointTransaction
+
+    with_witness do |w|
+      with_dummy_nested_transaction do |parent_tx|
+        top = parent_tx
+        parent_tx.transaction do |tx|
+          if tx.is_a?(DummyDriver::DummySavePointTransaction)
+            t = tx
+            w.check
+          end
+          raise DB::Rollback.new
+        end
+      end
+    end
+
+    t.rolledback.should be_true
+    t.committed.should be_false
+
+    top.rolledback.should be_false
+    top.committed.should be_true
+  end
+end

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