Extensions: Set Local & Migration Transaction Options (#26)

Add extensions: set_local, migration_transaction_options.

CHANGELOG.md

@@ -1,13 +1,26 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+## [0.7.0] 2022-06-24
+### Added
+- `DB.extension(:set_local)` - allows to set transaction locals;
+- Support of transaction options via `transaction_options` in migrations;
+
+## [0.6.0] 2022-06-15
+### Added
+- `mode` param for `Sequel::Model.plugin(:with_lock)`, defaults to `FOR NO KEY UPDATE`;
+
+## [0.5.0] 2020-06-06
+### Added
+- `Sequel::Model.plugin(:attr_encrypted)` - encrypts to model attributes;
+
 ## [0.4.0] 2019-11-18
 ### Added
 - `Sequel.extension(:deferrable_foreign_keys)` - makes foreign keys constraints deferrable by default;
 
 ## [0.3.2] 2018-07-03
 ### Added
-- Support sequel expessions in `with_rates`
+- Support sequel expressions in `with_rates`;
 
 ## [0.3.0] 2018-04-24
 ### Added

README.md

@@ -23,6 +23,8 @@ $ bundle
 - [`Synchronize`](#Synchronize)
 - [`Methods in Migrations`](#Methods-in-Migrations)
 - [`Deferrable Foreign Keys`](#Deferrable-Foreign-Keys)
+- [`Set Local`](#Set-Local)
+- [`Migration Transaction Options`](#Migration-Transaction-Options)
 
 # Plugins
 
@@ -204,6 +206,47 @@ end
 #   => <Husband @attributes={id:1, wife_id: 1}>
 ```
 
+
+## Set Local
+
+Enable: `DB.extension(:set_local)`
+
+Makes possible to set transaction locals.
+
+Example:
+
+```ruby
+DB.transaction(set_local: { lock_timeout: "5s", statement_timeout: "5s" }) {}
+```
+```sql
+BEGIN;
+SET LOCAL lock_timeout = '5s';
+SET LOCAL statement_timeout = '5s';
+COMMIT;
+```
+
+
+## Migration Transaction Options
+
+Enable: `Sequel.extension(:migration_transaction_options)`
+
+Makes possible to pass `transaction_options` in migrations.
+
+Example:
+
+```ruby
+Sequel.migration do
+  transaction_options rollback: :always
+
+  up { DB.select("1") }
+end
+```
+```sql
+BEGIN;
+SELECT '1';
+ROLLBACK;
+```
+
 ## AttrEncrypted
 
 Enable: `Sequel::Model.plugin :attr_encrypted`

lib/sequel/extensions/migration_transaction_options.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module MigrationDSLExtension
+  def transaction_options(opts)
+    migration.transaction_opts = opts
+  end
+end
+
+module SimpleMigrationExtension
+  attr_accessor :transaction_opts
+end
+
+module MigratorExtension
+  def checked_transaction(migration, &block)
+    if _use_transaction?(migration)
+      _transaction(migration, &block)
+    else
+      yield
+    end
+  end
+
+  private
+
+  def _use_transaction?(migration)
+    # NOTE: original code
+    if @use_transactions.nil?
+      if migration.use_transactions.nil?
+        @db.supports_transactional_ddl?
+      else
+        migration.use_transactions
+      end
+    else
+      @use_transactions
+    end
+  end
+
+  def _transaction(migration, &block)
+    if migration.transaction_opts.nil?
+      db.transaction(&block)
+    else
+      db.transaction(migration.transaction_opts, &block)
+    end
+  end
+end
+
+Sequel::MigrationDSL.include(MigrationDSLExtension)
+Sequel::SimpleMigration.include(SimpleMigrationExtension)
+Sequel::Migrator.prepend(MigratorExtension)

lib/sequel/extensions/set_local.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Sequel
+  module SetLocal
+    private
+
+    def begin_new_transaction(conn, opts)
+      super
+      check_set_local(conn, opts[:set_local])
+    end
+
+    def check_set_local(conn, locals)
+      return if locals.nil?
+
+      locals.each do |key, value|
+        log_connection_execute(conn, "SET LOCAL #{key} = \"#{value}\"")
+      end
+    end
+  end
+
+  Database.register_extension(:set_local, SetLocal)
+end

spec/extensions/migration_transaction_options_spec.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+RSpec.describe "migration_transaction_options" do
+  def migrator
+    Sequel::TimestampMigrator.new(DB, "spec/files/migrations")
+  end
+
+  before { migrator.run }
+
+  it "does not migrate since rollback always option is set" do
+    expect(DB.tables).to include(:first_table, :second_table)
+    expect(DB.tables).not_to include(:third_table)
+  end
+end

spec/extensions/set_local_spec.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+RSpec.describe "set_local" do
+  def run!
+    DB.transaction(set_local: { statement_timeout: "1s" }) do
+      DB.execute("SELECT pg_sleep(2)")
+    end
+  end
+
+  specify do
+    expect { run! }.to raise_error(/canceling statement due to statement timeout/)
+  end
+end

spec/files/migrations/1655939076_create_third_table.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+Sequel.migration do
+  transaction_options(rollback: :always)
+
+  up { create_table :third_table }
+  down { drop_table :third_table }
+end

utils/database.rb

@@ -14,6 +14,7 @@
 DB.extension :currency_rates
 DB.extension :pg_tools
 DB.extension :slave
+DB.extension :set_local
 DB.extension :synchronize
 
 Sequel.extension :deferrable_foreign_keys