Run specs in a transaction (#780)

This replaces truncating the database before every spec.
In avram we're seeing specs run ~20x faster but Lucky apps might expect to see speed ups of around 4-5x.

Author: matthewmcgarvey

db/migrations/20190723233131_test_change_type.cr

@@ -27,7 +27,7 @@ class TestChangeType::V20190723233131 < Avram::Migrator::Migration::V1
       change_type id : Int64
     end
 
-    TempUserInt64::BaseQuery.first # should not raise
+    TempUserInt64::BaseQuery.first.delete # should not raise
   end
 
   def rollback

spec/avram/database_spec.cr

@@ -1,7 +1,7 @@
 require "../spec_helper"
 
 describe Avram::Database do
-  describe "listen" do
+  describe "listen", tags: Avram::SpecHelper::TRUNCATE do
     it "yields the payload from a notify" do
       done = Channel(Nil).new
       TestDatabase.listen("dinner_time") do |notification|

spec/avram/instrumentation_spec.cr

@@ -2,7 +2,8 @@ require "../spec_helper"
 
 describe "Instrumentation" do
   it "publishes the query and args" do
-    TestDatabase.query "SELECT * FROM users"
+    # using block form to make sure the ResultSet is closed
+    TestDatabase.query "SELECT * FROM users" { }
 
     event = Avram::Events::QueryEvent.logged_events.last
     event.query.should eq("SELECT * FROM users")

spec/avram/model_spec.cr

@@ -33,7 +33,6 @@ describe Avram::Model do
     user = UserFactory.create
     post = PostFactory.create
 
-    user.id.should eq(post.id)
     user.should_not eq(post)
   end
 

spec/avram/query_logging_spec.cr

@@ -13,7 +13,7 @@ describe "Query logging" do
   it "does not log truncate statements" do
     Avram::QueryLog.dexter.temp_config do |log_io|
       TestDatabase.truncate
-      log_io.to_s.should eq("")
+      log_io.to_s.should_not contain("TRUNCATE TABLE")
     end
   end
 

spec/spec_helper.cr

@@ -19,8 +19,9 @@ Db::Create.new(quiet: true).run_task
 Db::Migrate.new(quiet: true).run_task
 Db::VerifyConnection.new(quiet: true).run_task
 
+Avram::SpecHelper.use_transactional_specs(TestDatabase)
+
 Spec.before_each do
-  TestDatabase.truncate
   # All specs seem to run on the same Fiber,
   # so we set back to NullStore before each spec
   # to ensure queries aren't randomly cached

src/avram.cr

@@ -7,6 +7,7 @@ require "db"
 require "pg"
 require "uuid"
 
+require "./ext/db/*"
 require "./avram/object_extensions"
 require "./avram/criteria"
 require "./avram/type"

src/avram/database.cr

@@ -3,7 +3,8 @@ abstract class Avram::Database
 
   @@db : DB::Database? = nil
   @@lock = Mutex.new
-  class_getter transactions = {} of FiberId => DB::Transaction
+  class_getter connections = {} of FiberId => DB::Connection
+  class_property lock_id : UInt64?
 
   macro inherited
     Habitat.create do
@@ -24,6 +25,16 @@ abstract class Avram::Database
     %}
   end
 
+  def self.setup_connection(&block : DB::Connection -> Nil)
+    new.db.setup_connection do |conn|
+      block.call conn
+    end
+  end
+
+  def self.verify_connection
+    new.connection.open.close
+  end
+
   # Rollback the current transaction
   def self.rollback
     new.rollback
@@ -142,28 +153,36 @@ abstract class Avram::Database
 
   # :nodoc:
   def run
-    yield current_transaction.try(&.connection) || db
+    yield current_connection || db
   end
 
   # :nodoc:
   def listen(*channels : String, &block : PQ::Notification ->) : Nil
     connection.connect_listen(*channels, &block)
   end
 
-  private def connection : Avram::Connection
+  protected def connection : Avram::Connection
     Avram::Connection.new(url, database_class: self.class)
   end
 
-  private def db : DB::Database
+  protected def db : DB::Database
     @@db ||= @@lock.synchronize do
       # check @@db again because a previous request could have set it after
       # the first time it was checked
       @@db || connection.open
     end
   end
 
+  private def current_connection : DB::Connection
+    connections[object_id] ||= db.checkout
+  end
+
+  private def object_id : UInt64
+    self.class.lock_id || Fiber.current.object_id
+  end
+
   private def current_transaction : DB::Transaction?
-    transactions[Fiber.current.object_id]?
+    current_connection._avram_stack.last?
   end
 
   protected def truncate
@@ -180,7 +199,7 @@ abstract class Avram::Database
 
   # :nodoc:
   def transaction : Bool
-    if current_transaction
+    if current_transaction.try(&._avram_joinable?)
       yield
       true
     else
@@ -190,20 +209,23 @@ abstract class Avram::Database
     end
   end
 
-  private def transactions
-    self.class.transactions
+  private def connections
+    self.class.connections
   end
 
   private def wrap_in_transaction
-    db.transaction do |tx|
-      transactions[Fiber.current.object_id] ||= tx
+    (current_transaction || current_connection).transaction do
       yield
     end
     true
   rescue e : Avram::Rollback
     false
   ensure
-    transactions.delete(Fiber.current.object_id)
+    # TODO: not sure of this
+    if current_connection._avram_in_transaction?
+      current_connection.release
+      connections.delete(object_id)
+    end
   end
 
   class DatabaseCleaner

src/avram/migrator/migration.cr

@@ -63,17 +63,17 @@ abstract class Avram::Migrator::Migration::V1
   end
 
   def migrated?
-    DB.open(Avram::Migrator::Runner.database_url) do |db|
-      db.query_one? "SELECT id FROM migrations WHERE version = $1", version, as: MigrationId
-    end
+    Avram.settings
+      .database_to_migrate
+      .query_one? "SELECT id FROM migrations WHERE version = $1", version, as: MigrationId
   end
 
-  private def track_migration(tx : DB::Transaction)
-    tx.connection.exec "INSERT INTO migrations(version) VALUES ($1)", version
+  private def track_migration(db : Avram::Database.class)
+    db.exec "INSERT INTO migrations(version) VALUES ($1)", version
   end
 
-  private def untrack_migration(tx : DB::Transaction)
-    tx.connection.exec "DELETE FROM migrations WHERE version = $1", version
+  private def untrack_migration(db : Avram::Database.class)
+    db.exec "DELETE FROM migrations WHERE version = $1", version
   end
 
   private def execute(statement : String)
@@ -87,16 +87,15 @@ abstract class Avram::Migrator::Migration::V1
   # # Usage
   #
   # ```
-  # execute_in_transaction ["DROP TABLE comments;"] do |tx|
-  #   tx.connection.exec "DROP TABLE users;"
+  # execute_in_transaction ["DROP TABLE comments;"] do |db|
+  #   db.exec "DROP TABLE users;"
   # end
   # ```
   private def execute_in_transaction(statements : Array(String))
-    DB.open(Avram::Migrator::Runner.database_url) do |db|
-      db.transaction do |tx|
-        statements.each { |s| tx.connection.exec s }
-        yield tx
-      end
+    database = Avram.settings.database_to_migrate
+    database.transaction do
+      statements.each { |s| database.exec s }
+      yield database
     end
   rescue e : PQ::PQError
     raise FailedMigration.new(migration: self.class.name, statements: statements, cause: e)

src/avram/migrator/runner.cr

@@ -11,6 +11,7 @@ class Avram::Migrator::Runner
   class_getter migrations = [] of Avram::Migrator::Migration::V1.class
 
   def initialize(@quiet : Bool = false)
+    Avram::Log.dexter.configure(:none)
   end
 
   def self.db_name
@@ -37,10 +38,6 @@ class Avram::Migrator::Runner
     Avram.settings.database_to_migrate.credentials
   end
 
-  def self.database_url
-    credentials.url
-  end
-
   def self.cmd_args
     String.build do |args|
       args << "-U #{self.db_user}" if self.db_user
@@ -106,9 +103,7 @@ class Avram::Migrator::Runner
   end
 
   def self.setup_migration_tracking_tables
-    DB.open(database_url) do |db|
-      db.exec create_table_for_tracking_migrations
-    end
+    Avram.settings.database_to_migrate.exec create_table_for_tracking_migrations
   end
 
   private def self.create_table_for_tracking_migrations

src/avram/spec_helper.cr

@@ -0,0 +1,55 @@
+module Avram::SpecHelper
+  TRUNCATE = "truncate"
+
+  macro use_transactional_specs(*databases)
+    Spec.around_each do |spec|
+      Avram::SpecHelper.wrap_spec_in_transaction(spec, {{ databases.splat }})
+    end
+  end
+
+  def self.wrap_spec_in_transaction(spec : Spec::Example::Procsy, *databases)
+    if use_truncation?(spec)
+      spec.run
+      databases.each(&.truncate)
+      return
+    end
+
+    tracked_transactions = [] of DB::Transaction
+
+    databases.each do |database|
+      database.lock_id = Fiber.current.object_id
+      database.connections.values.each do |conn|
+        tracked_transactions << conn.begin_transaction.tap(&._avram_joinable=(false))
+      end
+
+      database.setup_connection do |conn|
+        tracked_transactions << conn.begin_transaction.tap(&._avram_joinable=(false))
+      end
+    end
+
+    spec.run
+
+    tracked_transactions.each do |transaction|
+      next if transaction.closed? || transaction.connection.closed?
+
+      transaction.rollback
+      transaction.connection.release
+    end
+    tracked_transactions.clear
+    databases.each do |database|
+      database.connections.clear
+      database.setup_connection { }
+    end
+  end
+
+  private def self.use_truncation?(spec : Spec::Example::Procsy) : Bool
+    current = spec.example
+    while !current.is_a?(Spec::RootContext)
+      temp = current.as(Spec::Item)
+      return true if temp.tags.try(&.includes?(TRUNCATE))
+      current = temp.parent
+    end
+
+    false
+  end
+end

src/avram/tasks/db/verify_connection.cr

@@ -20,10 +20,9 @@ class Db::VerifyConnection < BaseTask
   end
 
   def run_task
-    DB.open(Avram::Migrator::Runner.database_url) do |_db|
-    end
+    Avram.settings.database_to_migrate.verify_connection
     puts "✔ Connection verified" unless quiet?
-  rescue PQ::ConnectionError | DB::ConnectionRefused
+  rescue Avram::ConnectionError
     raise <<-ERROR
     Unable to connect to Postgres for database '#{Avram.settings.database_to_migrate}'.
 

src/ext/db/connection.cr

@@ -0,0 +1,11 @@
+module DB
+  abstract class Connection
+    # :nodoc:
+    getter _avram_stack = [] of DB::Transaction
+
+    # :nodoc:
+    def _avram_in_transaction? : Bool
+      !_avram_stack.empty?
+    end
+  end
+end

src/ext/db/transaction.cr

@@ -0,0 +1,33 @@
+module DB
+  abstract class Transaction
+    # :nodoc:
+    property? _avram_joinable = true
+  end
+
+  class TopLevelTransaction < Transaction
+    def initialize(@connection : Connection)
+      @nested_transaction = false
+      @connection.perform_begin_transaction
+      @connection._avram_stack.push(self)
+    end
+
+    protected def do_close
+      @connection.release_from_transaction
+      @connection._avram_stack.pop
+    end
+  end
+
+  class SavePointTransaction < Transaction
+    def initialize(@parent : Transaction, @savepoint_name : String)
+      @nested_transaction = false
+      @connection = @parent.connection
+      @connection.perform_create_savepoint(@savepoint_name)
+      @connection._avram_stack.push(self)
+    end
+
+    protected def do_close
+      @parent.release_from_nested_transaction
+      @connection._avram_stack.pop
+    end
+  end
+end