activerecord-pg-extensions 0.2.3 → 0.4.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b19575bd4ea9eed35769514cf347333954d95087f57a65af0fc76f03e509ad96
-  data.tar.gz: 7c9e81b01d65797f3a47b344f89089f8d4f7788a5ae6e0ef2301165e370a6111
+  metadata.gz: 36ff139014c4b0b2288bf7839ff1f5d87ee596edb2f7f758f4eb66e2513b6479
+  data.tar.gz: 9880dd828e1d1d216e5797d8d94450b6f0fb41b01d32af31554f511f352c86c3
 SHA512:
-  metadata.gz: f32cdbdf6461d760ae956bd7cd7787fabff8d0b355c2c911e63e33c73357ce7170df4da05b1de5908612830e0bdaf754bba2cb8f9fe518d50d0064ccb5dfa5e9
-  data.tar.gz: 548ed3a2d3d90d98a68c5c637f009ae615e21199f764e88ed8ea954952a6ffa46036e022a5a514e5500f8d17ff22fe33080420209e38f9732cde67724a1218ed
+  metadata.gz: f7f15fbbf3b6475a894dc44d27380bdb141f25552d3c618550ab508e7e5d59d88b2d0dba297f40ca1d4ef030811942c00308674bdff89ac92a1a78a8c4d125a8
+  data.tar.gz: 69c17fc5576db6f1205b2a4819a87a26ac09fc3d98e43061e556cfb6b115d53e4ed326038632e2c2a219507e447e3e3d3199e879dbc97273c84227d25e2da27a

data/CHANGELOG.md

@@ -1,5 +1,9 @@
 ## [Unreleased]
 
+## [0.4.0] - 2021-09-29
+
+- Deprecate with_statement_timeout
+- Add statement_timeout, lock_timeout, and idle_in_transaction_session_timeout
 ## [0.2.3] - 2021-06-23
 
 - Fix extension_available?

data/lib/active_record/pg_extensions/errors.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  # @deprecated just use QueryCanceled
+  QueryTimeout = QueryCanceled
+end

data/lib/active_record/pg_extensions/pessimistic_migrations.rb

@@ -6,20 +6,36 @@ module ActiveRecord
     # to executing, in order to reduce the amount of time the actual DDL takes to
     # execute (and thus how long it needs the lock)
     module PessimisticMigrations
-      # does a query first to warm the db cache, to make the actual constraint adding fast
+      # adds a temporary check constraint to reduce locking when changing to NOT NULL, and we're not in a transaction
       def change_column_null(table, column, nullness, default = nil)
-        # no point in pre-warming the cache to avoid locking if we're already in a transaction
+        # no point in doing extra work to avoid locking if we're already in a transaction
         return super if nullness != false || open_transactions.positive?
+        return if columns(table).find { |c| c.name == column.to_s }&.null == false
+
+        temp_constraint_name = "chk_rails_#{table}_#{column}_not_null"
+        scope = quoted_scope(temp_constraint_name)
+        # check for temp constraint
+        valid = select_value(<<~SQL, "SCHEMA")
+          SELECT convalidated FROM pg_constraint INNER JOIN pg_namespace ON pg_namespace.oid=connamespace WHERE conname=#{scope[:name]} AND nspname=#{scope[:schema]}
+        SQL
+        if valid.nil?
+          add_check_constraint(table,
+                               "#{quote_column_name(column)} IS NOT NULL",
+                               name: temp_constraint_name,
+                               validate: false)
+        end
+        begin
+          validate_constraint(table, temp_constraint_name)
+        rescue ActiveRecord::StatementInvalid => e
+          raise ActiveRecord::NotNullViolation.new(sql: e.sql, binds: e.binds) if e.cause.is_a?(PG::CheckViolation)
+
+          raise
+        end
 
         transaction do
-          # make sure the query ignores indexes, because the actual ALTER TABLE will also ignore
-          # indexes
-          execute("SET LOCAL enable_indexscan=off")
-          execute("SET LOCAL enable_bitmapscan=off")
-          execute("SELECT COUNT(*) FROM #{quote_table_name(table)} WHERE #{quote_column_name(column)} IS NULL")
-          raise ActiveRecord::Rollback
+          super
+          remove_check_constraint(table, name: temp_constraint_name)
         end
-        super
       end
 
       # several improvements:

data/lib/active_record/pg_extensions/postgresql_adapter.rb

@@ -47,6 +47,7 @@ module ActiveRecord
         sql << " VERSION #{quote(version)}" if version
         sql << " CASCADE" if cascade
         execute(sql)
+        reload_type_map
         @extensions&.delete(extension.to_s)
       end
 
@@ -61,6 +62,7 @@ module ActiveRecord
         sql << " TO #{quote(version)}" if version && version != true
         sql << " SET SCHEMA #{schema}" if schema
         execute(sql)
+        reload_type_map
         @extensions&.delete(extension.to_s)
       end
 
@@ -73,6 +75,7 @@ module ActiveRecord
         sql << extensions.join(", ")
         sql << " CASCADE" if cascade
         execute(sql)
+        reload_type_map
         @extensions&.except!(*extensions.map(&:to_s))
       end
 
@@ -233,6 +236,56 @@ module ActiveRecord
         select_value("SELECT pg_is_in_recovery()")
       end
 
+      def set(configuration_parameter, value, local: false)
+        value = value.nil? ? "DEFAULT" : quote(value)
+        execute("SET#{' LOCAL' if local} #{configuration_parameter} TO #{value}")
+      end
+
+      def reset(configuration_parameter)
+        execute("RESET #{configuration_parameter}")
+      end
+
+      TIMEOUTS = %i[lock_timeout statement_timeout idle_in_transaction_session_timeout].freeze
+
+      TIMEOUTS.each do |kind|
+        class_eval <<~RUBY, __FILE__, __LINE__ + 1
+          def #{kind}
+            current_transaction.#{kind}
+          end
+
+          def #{kind}=(timeout)
+            raise ArgumentError, "Timeouts can only be set inside of a transaction" unless current_transaction.open?
+
+            current_transaction.send(:#{kind}=, timeout)
+          end
+        RUBY
+      end
+
+      # @deprecated: manage the transaction yourself and set statement_timeout directly
+      #
+      # otherwise, if you're already in a transaction, or you nest with_statement_timeout,
+      # the value will unexpectedly "stick" even after the block returns
+      def with_statement_timeout(timeout = nil)
+        timeout = 30 if timeout.nil? || timeout == true
+
+        transaction do
+          self.statement_timeout = timeout
+          yield
+        end
+      end
+
+      unless ::Rails.version >= "6.1"
+        def add_check_constraint(table_name, expression, name:, validate: true)
+          sql = +"ALTER TABLE #{quote_table_name(table_name)} ADD CONSTRAINT #{quote_column_name(name)} CHECK (#{expression})" # rubocop:disable Layout/LineLength
+          sql << " NOT VALID" unless validate
+          execute(sql)
+        end
+
+        def remove_check_constraint(table_name, name:)
+          execute("ALTER TABLE #{quote_table_name(table_name)} DROP CONSTRAINT #{quote_column_name(name)}")
+        end
+      end
+
       private
 
       if ::Rails.version < "6.1"

data/lib/active_record/pg_extensions/railtie.rb

@@ -8,9 +8,13 @@ module ActiveRecord
     class Railtie < Rails::Railtie
       initializer "pg_extensions.extend_ar", after: "active_record.initialize_database" do
         ActiveSupport.on_load(:active_record) do
+          require "active_record/pg_extensions/errors"
           require "active_record/pg_extensions/postgresql_adapter"
+          require "active_record/pg_extensions/transaction"
 
-          ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(PostgreSQLAdapter)
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(PostgreSQLAdapter)
+          ::ActiveRecord::ConnectionAdapters::NullTransaction.prepend(NullTransaction)
+          ::ActiveRecord::ConnectionAdapters::Transaction.prepend(Transaction)
           # if they've already require 'all', then inject now
           defined?(All) && All.inject
         end

data/lib/active_record/pg_extensions/transaction.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module PGExtensions
+    # Contains general additions to Transaction
+    module Transaction
+      PostgreSQLAdapter::TIMEOUTS.each do |kind|
+        class_eval <<~RUBY, __FILE__, __LINE__ + 1
+          def #{kind}(local: false)
+            return @#{kind} if local
+            @#{kind} || parent_transaction.#{kind}
+          end
+
+          private def #{kind}=(timeout)
+            return if @#{kind} == timeout
+
+            @#{kind} = timeout
+            return unless materialized?
+            connection.set(#{kind.inspect}, "\#{(timeout * 1000).to_i}ms", local: true)
+          end
+        RUBY
+      end
+
+      def materialize!
+        PostgreSQLAdapter::TIMEOUTS.each do |kind|
+          next if (timeout = send(kind, local: true)).nil?
+
+          connection.set(kind, "#{(timeout * 1000).to_i}ms", local: true)
+        end
+        super
+      end
+    end
+
+    # Contains general additions to NullTransaction
+    module NullTransaction
+      PostgreSQLAdapter::TIMEOUTS.each do |kind|
+        class_eval <<~RUBY, __FILE__, __LINE__ + 1
+          def #{kind}
+            nil
+          end
+        RUBY
+      end
+    end
+  end
+end

data/lib/active_record/pg_extensions/version.rb

@@ -2,6 +2,6 @@
 
 module ActiveRecord
   module PGExtensions
-    VERSION = "0.2.3"
+    VERSION = "0.4.2"
   end
 end

data/spec/pessimistic_migrations_spec.rb

@@ -6,32 +6,61 @@ describe ActiveRecord::PGExtensions::PessimisticMigrations do
   end
 
   describe "#change_column_null" do
+    # Rails 6.1 doesn't quote the constraint name when adding a check constraint??
+    def quote_constraint_name(name)
+      Rails.version >= "6.1" ? name : connection.quote_column_name(name)
+    end
+
     it "does nothing extra when changing a column to nullable" do
       connection.change_column_null(:table, :column, true)
       expect(connection.executed_statements).to eq ['ALTER TABLE "table" ALTER COLUMN "column" DROP NOT NULL']
     end
 
-    it "pre-warms the cache" do
+    it "does nothing if we're in a transaction" do
+      connection.transaction do
+        connection.change_column_null(:table, :column, true)
+      end
+      expect(connection.executed_statements).to eq ["BEGIN",
+                                                    'ALTER TABLE "table" ALTER COLUMN "column" DROP NOT NULL',
+                                                    "COMMIT"]
+    end
+
+    it "skips entirely if the column is already NOT NULL" do
+      expect(connection).to receive(:columns).with(:table).and_return([double(name: "column", null: false)])
       connection.change_column_null(:table, :column, false)
-      expect(connection.executed_statements).to eq(
-        ["BEGIN",
-         "SET LOCAL enable_indexscan=off",
-         "SET LOCAL enable_bitmapscan=off",
-         'SELECT COUNT(*) FROM "table" WHERE "column" IS NULL',
-         "ROLLBACK",
-         'ALTER TABLE "table" ALTER COLUMN "column" SET NOT NULL']
-      )
+      expect(connection.executed_statements).to eq([])
     end
 
-    it "does nothing extra if a transaction is already active" do
-      connection.transaction do
-        connection.change_column_null(:table, :column, false)
-      end
-      expect(connection.executed_statements).to eq(
-        ["BEGIN",
-         'ALTER TABLE "table" ALTER COLUMN "column" SET NOT NULL',
-         "COMMIT"]
-      )
+    it "adds and removes a check constraint" do
+      expect(connection).to receive(:columns).and_return([])
+      allow(connection).to receive(:check_constraint_for!).and_return(double(name: "chk_rails_table_column_not_null"))
+      connection.change_column_null(:table, :column, false)
+
+      expect(connection.executed_statements).to eq [
+        "SELECT convalidated FROM pg_constraint INNER JOIN pg_namespace ON pg_namespace.oid=connamespace WHERE conname='chk_rails_table_column_not_null' AND nspname=ANY (current_schemas(false))\n", # rubocop:disable Layout/LineLength
+        %{ALTER TABLE "table" ADD CONSTRAINT #{quote_constraint_name('chk_rails_table_column_not_null')} CHECK ("column" IS NOT NULL) NOT VALID}, # rubocop:disable Layout/LineLength
+        'ALTER TABLE "table" VALIDATE CONSTRAINT "chk_rails_table_column_not_null"',
+        "BEGIN",
+        'ALTER TABLE "table" ALTER COLUMN "column" SET NOT NULL',
+        'ALTER TABLE "table" DROP CONSTRAINT "chk_rails_table_column_not_null"',
+        "COMMIT"
+      ]
+    end
+
+    it "verifies an existing check constraint" do
+      expect(connection).to receive(:columns).and_return([])
+      allow(connection).to receive(:check_constraint_for!).and_return(double(name: "chk_rails_table_column_not_null"))
+      expect(connection).to receive(:select_value).and_return(false)
+      connection.change_column_null(:table, :column, false)
+
+      expect(connection.executed_statements).to eq [
+        # stubbed out <check constraint valid>
+        'ALTER TABLE "table" VALIDATE CONSTRAINT "chk_rails_table_column_not_null"',
+        "BEGIN",
+        'ALTER TABLE "table" ALTER COLUMN "column" SET NOT NULL',
+        'ALTER TABLE "table" DROP CONSTRAINT "chk_rails_table_column_not_null"',
+        "COMMIT"
+      ]
     end
   end
 

data/spec/postgresql_adapter_spec.rb

@@ -1,6 +1,10 @@
 # frozen_string_literal: true
 
 describe ActiveRecord::ConnectionAdapters::PostgreSQLAdapter do
+  before do
+    allow(connection).to receive(:reload_type_map)
+  end
+
   describe "#set_constraints" do
     around do |example|
       connection.dont_execute(&example)
@@ -264,4 +268,190 @@ describe ActiveRecord::ConnectionAdapters::PostgreSQLAdapter do
       end
     end
   end
+
+  describe "#with_statement_timeout" do
+    around do |example|
+      # these specs were written before we supported deferring setting timeouts
+      # until the transaction materializes
+      connection.disable_lazy_transactions!
+      example.call
+      connection.enable_lazy_transactions!
+    end
+
+    it "stops long-running queries" do
+      expect do
+        connection.with_statement_timeout(0.01) do
+          connection.execute("SELECT pg_sleep(3)")
+        end
+      end.to raise_error(ActiveRecord::QueryCanceled)
+    end
+
+    it "re-raises other errors" do
+      expect do
+        connection.with_statement_timeout(1) do
+          connection.execute("bad sql")
+        end
+      end.to(raise_error { |e| expect(e.cause).to be_a(PG::SyntaxError) })
+    end
+
+    context "without executing" do
+      around do |example|
+        connection.dont_execute(&example)
+      end
+
+      it "converts integer to ms" do
+        connection.with_statement_timeout(30) { nil }
+        expect(connection.executed_statements).to eq(
+          [
+            "BEGIN",
+            "SET LOCAL statement_timeout TO '30000ms'",
+            "COMMIT"
+          ]
+        )
+      end
+
+      it "converts float to ms" do
+        connection.with_statement_timeout(5.5) { nil }
+        expect(connection.executed_statements).to eq(
+          [
+            "BEGIN",
+            "SET LOCAL statement_timeout TO '5500ms'",
+            "COMMIT"
+          ]
+        )
+      end
+
+      it "converts ActiveSupport::Duration to ms" do
+        connection.with_statement_timeout(5.seconds) { nil }
+        expect(connection.executed_statements).to eq(
+          [
+            "BEGIN",
+            "SET LOCAL statement_timeout TO '5000ms'",
+            "COMMIT"
+          ]
+        )
+      end
+
+      it "allows true" do
+        connection.with_statement_timeout(true) { nil }
+        expect(connection.executed_statements).to eq(
+          [
+            "BEGIN",
+            "SET LOCAL statement_timeout TO '30000ms'",
+            "COMMIT"
+          ]
+        )
+      end
+    end
+  end
+
+  describe "#statement_timeout=" do
+    around do |example|
+      connection.dont_execute(&example)
+    end
+
+    it "raises if a transaction isn't active" do
+      expect { connection.statement_timeout = 30 }.to raise_error(ArgumentError)
+    end
+
+    it "does nothing if the transaction never materializes" do
+      connection.transaction do
+        connection.statement_timeout = 30
+        expect(connection.statement_timeout).to eq 30
+      end
+      expect(connection.statement_timeout).to be_nil
+
+      expect(connection.executed_statements).to be_empty
+    end
+
+    it "sets the timeout if the transaction is materialized" do
+      connection.transaction do
+        connection.select_value("SELECT 1")
+        connection.statement_timeout = 30
+        expect(connection.statement_timeout).to eq 30
+      end
+      expect(connection.statement_timeout).to be_nil
+
+      expect(connection.executed_statements).to eq(
+        ["BEGIN",
+         "SELECT 1",
+         "SET LOCAL statement_timeout TO '30000ms'",
+         "COMMIT"]
+      )
+    end
+
+    it "sets the timeout if the transaction materializes" do
+      connection.transaction do
+        connection.statement_timeout = 30
+        connection.select_value("SELECT 1")
+        expect(connection.statement_timeout).to eq 30
+      end
+      expect(connection.statement_timeout).to be_nil
+
+      expect(connection.executed_statements).to eq(
+        ["BEGIN",
+         "SET LOCAL statement_timeout TO '30000ms'",
+         "SELECT 1",
+         "COMMIT"]
+      )
+    end
+
+    it "works with nested transactions" do
+      connection.transaction do
+        connection.statement_timeout = 30
+        connection.transaction(requires_new: true) do
+          connection.statement_timeout = 15
+          connection.select_value("SELECT 1")
+          expect(connection.statement_timeout).to eq 15
+        end
+        expect(connection.statement_timeout).to eq 30
+      end
+      expect(connection.statement_timeout).to be_nil
+
+      expect(connection.executed_statements).to eq(
+        ["BEGIN",
+         "SET LOCAL statement_timeout TO '30000ms'",
+         "SAVEPOINT active_record_1",
+         "SET LOCAL statement_timeout TO '15000ms'",
+         "SELECT 1",
+         "RELEASE SAVEPOINT active_record_1",
+         "COMMIT"]
+      )
+    end
+  end
+
+  unless Rails.version >= "6.1"
+    describe "#add_check_constraint" do
+      around do |example|
+        connection.dont_execute(&example)
+      end
+
+      it "works" do
+        connection.add_check_constraint(:table, "column IS NOT NULL", name: :my_constraint)
+        expect(connection.executed_statements).to eq(
+          ['ALTER TABLE "table" ADD CONSTRAINT "my_constraint" CHECK (column IS NOT NULL)']
+        )
+      end
+
+      it "defers validation" do
+        connection.add_check_constraint(:table, "column IS NOT NULL", name: :my_constraint, validate: false)
+        expect(connection.executed_statements).to eq(
+          ['ALTER TABLE "table" ADD CONSTRAINT "my_constraint" CHECK (column IS NOT NULL) NOT VALID']
+        )
+      end
+    end
+
+    describe "#remove_check_constraint" do
+      around do |example|
+        connection.dont_execute(&example)
+      end
+
+      it "works" do
+        connection.remove_check_constraint(:table, name: :my_constraint)
+        expect(connection.executed_statements).to eq(
+          ['ALTER TABLE "table" DROP CONSTRAINT "my_constraint"']
+        )
+      end
+    end
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-pg-extensions
 version: !ruby/object:Gem::Version
-  version: 0.2.3
+  version: 0.4.2
 platform: ruby
 authors:
 - Cody Cutrer
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-06-22 00:00:00.000000000 Z
+date: 2021-10-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -175,10 +175,12 @@ files:
 - config/database.yml
 - lib/active_record/pg_extensions.rb
 - lib/active_record/pg_extensions/all.rb
+- lib/active_record/pg_extensions/errors.rb
 - lib/active_record/pg_extensions/extension.rb
 - lib/active_record/pg_extensions/pessimistic_migrations.rb
 - lib/active_record/pg_extensions/postgresql_adapter.rb
 - lib/active_record/pg_extensions/railtie.rb
+- lib/active_record/pg_extensions/transaction.rb
 - lib/active_record/pg_extensions/version.rb
 - lib/activerecord-pg-extensions.rb
 - spec/pessimistic_migrations_spec.rb
@@ -204,7 +206,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3
+rubygems_version: 3.2.24
 signing_key:
 specification_version: 4
 summary: Several extensions to ActiveRecord's PostgreSQLAdapter.