activerecord-pg-extensions 0.2.3 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b19575bd4ea9eed35769514cf347333954d95087f57a65af0fc76f03e509ad96
-  data.tar.gz: 7c9e81b01d65797f3a47b344f89089f8d4f7788a5ae6e0ef2301165e370a6111
+  metadata.gz: 3b240f7e9306ba57d7f52add8cef37ebe735e2f3d900d328d7cfb2e22fe1fcc6
+  data.tar.gz: 04b62ecf9b4c250e2f664f1a0be0c71baac1b6f5d117dc8d45f08bb7b7100459
 SHA512:
-  metadata.gz: f32cdbdf6461d760ae956bd7cd7787fabff8d0b355c2c911e63e33c73357ce7170df4da05b1de5908612830e0bdaf754bba2cb8f9fe518d50d0064ccb5dfa5e9
-  data.tar.gz: 548ed3a2d3d90d98a68c5c637f009ae615e21199f764e88ed8ea954952a6ffa46036e022a5a514e5500f8d17ff22fe33080420209e38f9732cde67724a1218ed
+  metadata.gz: c5a1ae9925eb093ea81424574d409186c85ef2474b1cc2a38f388757ff4e94a8b3893ca7e582e5010e8dbe8915a2e499d1600103c12936169d0447703b5c8d68
+  data.tar.gz: 256b8740a417e1bcc84a7edd23663da83543aa412a7221ae2a818721dd04a16b7ccf3f06915cb625d50f95ecb55ccd0719e1d8204441294dd9860ba25cb54b6a

data/lib/active_record/pg_extensions/errors.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class QueryTimeout < ActiveRecord::StatementInvalid; end
+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,30 @@ module ActiveRecord
         select_value("SELECT pg_is_in_recovery()")
       end
 
+      def with_statement_timeout(timeout = nil)
+        timeout = 30 if timeout.nil? || timeout == true
+        transaction do
+          execute("SET LOCAL statement_timeout=#{(timeout * 1000).to_i}")
+          yield
+        rescue ActiveRecord::StatementInvalid => e
+          raise ActiveRecord::QueryTimeout.new(sql: e.sql, binds: e.binds) if e.cause.is_a?(PG::QueryCanceled)
+
+          raise
+        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,6 +8,7 @@ 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"
 
           ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(PostgreSQLAdapter)

data/lib/active_record/pg_extensions/version.rb

@@ -2,6 +2,6 @@
 
 module ActiveRecord
   module PGExtensions
-    VERSION = "0.2.3"
+    VERSION = "0.3.0"
   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,103 @@ describe ActiveRecord::ConnectionAdapters::PostgreSQLAdapter do
       end
     end
   end
+
+  describe "#with_statement_timeout" do
+    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::QueryTimeout)
+    end
+
+    it "re-raises other errors" do
+      expect do
+        connection.with_statement_timeout(1) do
+          connection.execute("bad sql")
+        end
+      end.to raise_error(ActiveRecord::StatementInvalid)
+    end
+
+    context "without executing" do
+      it "converts integer to ms" do
+        connection.with_statement_timeout(30) { nil }
+        expect(connection.executed_statements).to eq(
+          [
+            "BEGIN",
+            "SET LOCAL statement_timeout=30000",
+            "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=5500",
+            "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=5000",
+            "COMMIT"
+          ]
+        )
+      end
+
+      it "allows true" do
+        connection.with_statement_timeout(true) { nil }
+        expect(connection.executed_statements).to eq(
+          [
+            "BEGIN",
+            "SET LOCAL statement_timeout=30000",
+            "COMMIT"
+          ]
+        )
+      end
+    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.3.0
 platform: ruby
 authors:
 - Cody Cutrer
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-06-22 00:00:00.000000000 Z
+date: 2021-09-09 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -175,6 +175,7 @@ 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
@@ -204,7 +205,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.