activerecord-pg-extensions 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1c73159dc2062391844fc7309dd066720684b43a7d81c4c19c8bc5acaeacb36d
-  data.tar.gz: bff7754ec2e0a7817b7806f0e37c80ec91675a2e745841d60421da9192438840
+  metadata.gz: abfb428bb45f993d1a1ee897e81be0372dafc62b5a345c433eb9e69c821ba100
+  data.tar.gz: d8a2f84e3a2a70576d165e1ad56a17d2cb122acc1bf784a9a17c969cc3f018a8
 SHA512:
-  metadata.gz: a8ed2044c44c79a5d6f0412ebb95de80e9a43361d22866950605a14c2a59bcf9566a6cd73febe06743d5e4f8c445a76b3c25c9ed7395733ba9159a386e1bbd2b
-  data.tar.gz: 3c2b19dccfa6b809f5e9de4a0c94c5bb26000537175957a2782bf8e94daf7f568dd311c45c36261f843128b53f8f9ff038c59a9141a952104316c7cb5b91f4d2
+  metadata.gz: 4e83d36db41c570b469ccac0e931b0a3ce04536319b9a357e28225b6c3adb613d1cd300993b7eb9c45733ea54e2b3c3b2ab79a51a932c428aa054dae13fb138d
+  data.tar.gz: 0f404b7d50dd5e48cdf729987888c164f133301da6b1be5867fa6d8c1aa3ba5223eae556778c082ad31afb10bccdc2eecba1f166f69c3b12311630047dea2a55

data/CHANGELOG.md

@@ -1,5 +1,13 @@
 ## [Unreleased]
 
+## [0.2.0] - 2021-06-07
+
+- Add PostgreSQLAdapter#set_replica_identity
+- Add methods for discovering and managing extensions
+- Add method to temporarily add a schema to the search path
+- Add PostgreSQLAdapter#vacuum
+- Add optional module PessimisticMigrations
+
 ## [0.1.1] - 2021-06-07
 
 - Add PostgreSQLAdapter#defer_constraints

data/lib/active_record/pg_extensions/all.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require "active_record/pg_extensions/pessimistic_migrations"
+
+module ActiveRecord
+  module PGExtensions
+    # includes all optional extensions at once
+    module All
+      def self.inject
+        ConnectionAdapters::PostgreSQLAdapter.prepend(PessimisticMigrations)
+      end
+    end
+  end
+end
+
+ActiveRecord::PGExtensions::All.inject if defined?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)

data/lib/active_record/pg_extensions/extension.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module PGExtensions
+    # Contains general additions to the PostgreSQLAdapter
+    module PostgreSQLAdapter
+      Extension = Struct.new(:name, :schema, :version)
+    end
+  end
+end

data/lib/active_record/pg_extensions/pessimistic_migrations.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module PGExtensions
+    # Changes several DDL commands to trigger background queries to warm caches prior
+    # 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
+      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
+        return super if nullness != false || open_transactions.positive?
+
+        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
+        end
+        super
+      end
+
+      # several improvements:
+      #  * support if_not_exists
+      #  * delay_validation automatically creates the FK as NOT VALID, and then immediately validates it
+      #  * if delay_validation is used, and the index already exists but is NOT VALID, it just re-tries
+      #    the validation, instead of failing
+      def add_foreign_key(from_table, to_table, delay_validation: false, if_not_exists: false, **options)
+        # pointless if we're in a transaction
+        delay_validation = false if open_transactions.positive?
+        options[:validate] = false if delay_validation
+
+        options = foreign_key_options(from_table, to_table, options)
+
+        if if_not_exists || delay_validation
+          scope = quoted_scope(options[:name])
+          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
+          return if valid == true && if_not_exists
+        end
+
+        super(from_table, to_table, **options) unless valid == false
+        validate_constraint(from_table, options[:name]) if delay_validation
+      end
+
+      # will automatically remove a NOT VALID index before trying to add
+      def add_index(table_name, column_name, options = {})
+        # catch a concurrent index add that fails because it already exists, and is invalid
+        if options[:algorithm] == :concurrently || options[:if_not_exists]
+          column_names = index_column_names(column_name)
+          index_name = options[:name].to_s if options.key?(:name)
+          index_name ||= index_name(table_name, column_names)
+
+          index = quoted_scope(index_name)
+          table = quoted_scope(table_name)
+          valid = select_value(<<~SQL, "SCHEMA")
+            SELECT indisvalid
+            FROM pg_class t
+            INNER JOIN pg_index d ON t.oid = d.indrelid
+            INNER JOIN pg_class i ON d.indexrelid = i.oid
+            WHERE i.relkind = 'i'
+              AND i.relname = #{index[:name]}
+              AND t.relname = #{table[:name]}
+              AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = #{index[:schema]} )
+            LIMIT 1
+          SQL
+          if valid == false && options[:algorithm] == :concurrently
+            remove_index(table_name, name: index_name,
+                                     algorithm: :concurrently)
+          end
+          return if options[:if_not_exists] && valid == true
+        end
+        # Rails.version: can stop doing this in Rails 6.2, when it's natively supported
+        options.delete(:if_not_exists)
+        super
+      end
+    end
+  end
+end

data/lib/active_record/pg_extensions/postgresql_adapter.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "active_record/pg_extensions/extension"
+
 module ActiveRecord
   module PGExtensions
     # Contains general additions to the PostgreSQLAdapter
@@ -24,6 +26,226 @@ module ActiveRecord
         yield
         set_constraints(:immediate, *constraints)
       end
+
+      # see https://www.postgresql.org/docs/current/sql-altertable.html#SQL-CREATETABLE-REPLICA-IDENTITY
+      def set_replica_identity(table, identity = :default)
+        identity_clause = case identity
+                          when :default, :full, :nothing
+                            identity.to_s.upcase
+                          else
+                            "USING INDEX #{quote_column_name(identity)}"
+                          end
+        execute("ALTER TABLE #{quote_table_name(table)} REPLICA IDENTITY #{identity_clause}")
+      end
+
+      # see https://www.postgresql.org/docs/current/sql-createextension.html
+      def create_extension(extension, if_not_exists: false, schema: nil, version: nil, cascade: false)
+        sql = +"CREATE EXTENSION "
+        sql <<= "IF NOT EXISTS " if if_not_exists
+        sql << extension.to_s
+        sql << " SCHEMA #{schema}" if schema
+        sql << " VERSION #{quote(version)}" if version
+        sql << " CASCADE" if cascade
+        execute(sql)
+        @extensions&.delete(extension.to_s)
+      end
+
+      # see https://www.postgresql.org/docs/current/sql-alterextension.html
+      def alter_extension(extension, schema: nil, version: nil)
+        if schema && version
+          raise ArgumentError, "Cannot change schema and upgrade to a particular version in a single statement"
+        end
+
+        sql = +"ALTER EXTENSION #{extension}"
+        sql << " UPDATE" if version
+        sql << " TO #{quote(version)}" if version && version != true
+        sql << " SET SCHEMA #{schema}" if schema
+        execute(sql)
+        @extensions&.delete(extension.to_s)
+      end
+
+      # see https://www.postgresql.org/docs/current/sql-dropextension.html
+      def drop_extension(*extensions, if_exists: false, cascade: false)
+        raise ArgumentError, "wrong number of arguments (given 0, expected 1+)" if extensions.empty?
+
+        sql = +"DROP EXTENSION "
+        sql << "IF EXISTS " if if_exists
+        sql << extensions.join(", ")
+        sql << " CASCADE" if cascade
+        execute(sql)
+        @extensions&.except!(*extensions.map(&:to_s))
+      end
+
+      # check if a particular extension can be installed
+      def extension_available?(extension, version = nil)
+        sql = +"SELECT 1 FROM "
+        sql << version ? "pg_available_extensions" : "pg_available_extension_versions"
+        sql << " WHERE name=#{quote(extension)}"
+        sql << " AND version=#{quote(version)}" if version
+        select_value(sql).to_i == 1
+      end
+
+      # returns an Extension object for a particular extension
+      def extension(extension)
+        @extensions ||= {}
+        @extensions.fetch(extension.to_s) do
+          rows = select_rows(<<~SQL, "SCHEMA")
+            SELECT nspname, extversion
+            FROM pg_extension
+              INNER JOIN pg_namespace ON extnamespace=pg_namespace.oid
+            WHERE extname=#{quote(extension)}
+          SQL
+          next nil if rows.empty?
+
+          Extension.new(extension.to_s, rows[0][0], rows[0][1])
+        end
+      end
+
+      # temporarily adds schema to the search_path (i.e. so you can use an extension that won't work
+      # without being on the search path, such as postgis)
+      def add_schema_to_search_path(schema)
+        if schema_search_path.split(",").include?(schema)
+          yield
+        else
+          old_search_path = schema_search_path
+          manual_rollback = false
+          transaction(requires_new: true) do
+            self.schema_search_path += ",#{schema}"
+            yield
+            self.schema_search_path = old_search_path
+          rescue ActiveRecord::StatementInvalid, ActiveRecord::Rollback => e
+            # the transaction rolling back will revert the search path change;
+            # we don't need to do another query to set it
+            @schema_search_path = old_search_path
+            manual_rollback = e if e.is_a?(ActiveRecord::Rollback)
+            raise
+          end
+          # the transaction call will swallow ActiveRecord::Rollback,
+          # but we want it this method to be transparent
+          raise manual_rollback if manual_rollback
+        end
+      end
+
+      # see https://www.postgresql.org/docs/current/sql-vacuum.html
+      def vacuum(*table_and_columns,
+                 full: false,
+                 freeze: false,
+                 verbose: false,
+                 analyze: false,
+                 disable_page_skipping: false,
+                 skip_locked: false,
+                 index_cleanup: false,
+                 truncate: false,
+                 parallel: nil)
+        if parallel && !(parallel.is_a?(Integer) && parallel.positive?)
+          raise ArgumentError, "parallel must be a positive integer"
+        end
+
+        sql = +"VACUUM"
+        sql << " FULL" if full
+        sql << " FREEZE" if freeze
+        sql << " VERBOSE" if verbose
+        sql << " ANALYZE" if analyze
+        sql << " DISABLE_PAGE_SKIPPING" if disable_page_skipping
+        sql << " SKIP_LOCKED" if skip_locked
+        sql << " INDEX_CLEANUP" if index_cleanup
+        sql << " TRUNCATE" if truncate
+        sql << " PARALLEL #{parallel}" if parallel
+        sql << " " unless table_and_columns.empty?
+        sql << table_and_columns.map do |table|
+          if table.is_a?(Hash)
+            raise ArgumentError, "columns may only be specified if a analyze is specified" unless analyze
+
+            table.map do |table_name, columns|
+              "#{quote_table_name(table_name)} (#{Array.wrap(columns).map { |c| quote_column_name(c) }.join(', ')})"
+            end.join(", ")
+          else
+            quote_table_name(table)
+          end
+        end.join(", ")
+        execute(sql)
+      end
+
+      # Amazon Aurora doesn't have a WAL
+      def wal?
+        unless instance_variable_defined?(:@has_wal)
+          function_name = pre_pg10_wal_function_name("pg_current_wal_lsn")
+          @has_wal = select_value("SELECT true FROM pg_proc WHERE proname='#{function_name}' LIMIT 1")
+        end
+        @has_wal
+      end
+
+      # see https://www.postgresql.org/docs/current/functions-admin.html#id-1.5.8.33.5.5.2.2.4.1.1.1
+      def current_wal_lsn
+        return nil unless wal?
+
+        select_value("SELECT #{pre_pg10_wal_function_name('pg_current_wal_lsn')}()")
+      end
+
+      # see https://www.postgresql.org/docs/current/functions-admin.html#id-1.5.8.33.5.5.2.2.2.1.1.1
+      def current_wal_flush_lsn
+        return nil unless wal?
+
+        select_value("SELECT #{pre_pg10_wal_function_name('pg_current_wal_flush_lsn')}()")
+      end
+
+      # see https://www.postgresql.org/docs/current/functions-admin.html#id-1.5.8.33.5.5.2.2.3.1.1.1
+      def current_wal_insert_lsn
+        return nil unless wal?
+
+        select_value("SELECT #{pre_pg10_wal_function_name('pg_current_wal_insert_lsn')}()")
+      end
+
+      # https://www.postgresql.org/docs/current/functions-admin.html#id-1.5.8.33.6.3.2.2.2.1.1.1
+      def last_wal_receive_lsn
+        return nil unless wal?
+
+        select_value("SELECT #{pre_pg10_wal_function_name('pg_last_wal_receive_lsn')}()")
+      end
+
+      # see https://www.postgresql.org/docs/current/functions-admin.html#id-1.5.8.33.6.3.2.2.3.1.1.1
+      def last_wal_replay_lsn
+        return nil unless wal?
+
+        select_value("SELECT #{pre_pg10_wal_function_name('pg_last_wal_replay_lsn')}()")
+      end
+
+      # see https://www.postgresql.org/docs/current/functions-admin.html#id-1.5.8.33.5.5.2.2.4.1.1.1
+      # lsns can be literals, or :current, :current_flush, :current_insert, :last_receive, or :last_replay
+      def wal_lsn_diff(lsn1 = :current, lsn2 = :last_replay)
+        return nil unless wal?
+
+        lsns = [lsn1, lsn2].map do |lsn|
+          case lsn
+          when :current then pre_pg10_wal_function_name("pg_current_wal_lsn()")
+          when :current_flush then pre_pg10_wal_function_name("pg_current_flush_wal_lsn()")
+          when :current_insert then pre_pg10_wal_function_name("pg_current_insert_wal_lsn()")
+          when :last_receive then pre_pg10_wal_function_name("pg_last_wal_receive_lsn()")
+          when :last_replay then pre_pg10_wal_function_name("pg_last_wal_replay_lsn()")
+          else; quote(lsn)
+          end
+        end
+
+        select_value("SELECT #{pre_pg10_wal_function_name('pg_wal_lsn_diff')}(#{lsns[0]}, #{lsns[1]})")
+      end
+
+      def in_recovery?
+        select_value("SELECT pg_is_in_recovery()")
+      end
+
+      private
+
+      def initialize_type_map(map = type_map)
+        map.register_type "pg_lsn", ActiveRecord::ConnectionAdapters::PostgreSQL::OID::SpecializedString.new(:pg_lsn)
+
+        super
+      end
+
+      def pre_pg10_wal_function_name(func)
+        return func if postgresql_version >= 100_000
+
+        func.sub("wal", "xlog").sub("lsn", "location")
+      end
     end
   end
 end

data/lib/active_record/pg_extensions/railtie.rb

@@ -11,6 +11,8 @@ module ActiveRecord
           require "active_record/pg_extensions/postgresql_adapter"
 
           ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.include(PostgreSQLAdapter)
+          # if they've already require 'all', then inject now
+          defined?(All) && All.inject
         end
       end
     end

data/lib/active_record/pg_extensions/version.rb

@@ -2,6 +2,6 @@
 
 module ActiveRecord
   module PGExtensions
-    VERSION = "0.1.1"
+    VERSION = "0.2.0"
   end
 end

data/spec/pessimistic_migrations_spec.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+describe ActiveRecord::PGExtensions::PessimisticMigrations do
+  around do |example|
+    connection.dont_execute(&example)
+  end
+
+  describe "#change_column_null" do
+    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
+      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']
+      )
+    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"]
+      )
+    end
+  end
+
+  describe "#add_foreign_key" do
+    it "does nothing extra if a transaction is already active" do
+      connection.transaction do
+        connection.add_foreign_key :emails, :users, delay_validation: true
+      end
+      expect(connection.executed_statements).to match(
+        ["BEGIN",
+         match(/\AALTER TABLE "emails" ADD CONSTRAINT "fk_rails_[0-9a-f]+".+REFERENCES "users" \("id"\)\s*\z/m),
+         "COMMIT"]
+      )
+    end
+
+    it "delays validation" do
+      connection.add_foreign_key :emails, :users, delay_validation: true
+      expect(connection.executed_statements).to match(
+        [/convalidated/,
+         match(/\AALTER TABLE "emails" ADD CONSTRAINT "[a-z0-9_]+".+REFERENCES "users" \("id"\)\s+NOT VALID\z/m),
+         match(/^ALTER TABLE "emails" VALIDATE CONSTRAINT "fk_rails_[0-9a-f]+"/)]
+      )
+    end
+
+    it "only validates if the constraint already exists, and is not valid" do
+      expect(connection).to receive(:select_value).with(/convalidated/, "SCHEMA").and_return(false)
+      connection.add_foreign_key :emails, :users, delay_validation: true
+      expect(connection.executed_statements).to match(
+        [match(/^ALTER TABLE "emails" VALIDATE CONSTRAINT "fk_rails_[0-9a-f]+"/)]
+      )
+    end
+
+    it "does nothing if constraint already exists" do
+      expect(connection).to receive(:select_value).with(/convalidated/, "SCHEMA").and_return(true)
+      connection.add_foreign_key :emails, :users, if_not_exists: true
+      expect(connection.executed_statements).to eq []
+    end
+
+    it "still tries if delay_validation is true but if_not_exists is false and it already exists" do
+      expect(connection).to receive(:select_value).with(/convalidated/, "SCHEMA").and_return(true)
+      connection.add_foreign_key :emails, :users, delay_validation: true
+      expect(connection.executed_statements).to match(
+        [match(/\AALTER TABLE "emails" ADD CONSTRAINT "[a-z0-9_]+".+REFERENCES "users" \("id"\)\s+NOT VALID\z/m),
+         match(/^ALTER TABLE "emails" VALIDATE CONSTRAINT "fk_rails_[0-9a-f]+"/)]
+      )
+    end
+
+    it "does nothing if_not_exists is true and it is NOT VALID" do
+      expect(connection).to receive(:select_value).with(/convalidated/, "SCHEMA").and_return(false)
+      connection.add_foreign_key :emails, :users, if_not_exists: true
+      expect(connection.executed_statements).to eq []
+    end
+  end
+
+  describe "#add_index" do
+    it "removes a NOT VALID index before re-adding" do
+      expect(connection).to receive(:select_value).with(/indisvalid/, "SCHEMA").and_return(false)
+      expect(connection).to receive(:remove_index).with(:users, name: "index_users_on_name", algorithm: :concurrently)
+
+      connection.add_index :users, :name, algorithm: :concurrently
+      expect(connection.executed_statements).to eq(
+        ['CREATE INDEX CONCURRENTLY "index_users_on_name" ON "users" ("name")']
+      )
+    end
+
+    it "does nothing if the index already exists" do
+      expect(connection).to receive(:select_value).with(/indisvalid/, "SCHEMA").and_return(true)
+
+      connection.add_index :users, :name, if_not_exists: true
+      expect(connection.executed_statements).to eq []
+    end
+  end
+end

data/spec/postgresql_adapter_spec.rb

@@ -40,4 +40,208 @@ describe ActiveRecord::ConnectionAdapters::PostgreSQLAdapter do
       expect(connection.executed_statements).to eq ["SET CONSTRAINTS ALL DEFERRED", "SET CONSTRAINTS ALL IMMEDIATE"]
     end
   end
+
+  describe "#set_replica_identity" do
+    around do |example|
+      connection.dont_execute(&example)
+    end
+
+    it "resets identity" do
+      connection.set_replica_identity(:my_table)
+      expect(connection.executed_statements).to eq ['ALTER TABLE "my_table" REPLICA IDENTITY DEFAULT']
+    end
+
+    it "sets full identity" do
+      connection.set_replica_identity(:my_table, :full)
+      expect(connection.executed_statements).to eq ['ALTER TABLE "my_table" REPLICA IDENTITY FULL']
+    end
+
+    it "sets an index" do
+      connection.set_replica_identity(:my_table, :my_index)
+      expect(connection.executed_statements).to eq ['ALTER TABLE "my_table" REPLICA IDENTITY USING INDEX "my_index"']
+    end
+  end
+
+  context "extensions" do
+    it "creates and drops an extension" do
+      connection.create_extension(:pg_trgm, schema: "public")
+      expect(connection.executed_statements).to eq ["CREATE EXTENSION pg_trgm SCHEMA public"]
+      expect(ext = connection.extension(:pg_trgm)).not_to be_nil
+      expect(ext.schema).to eq "public"
+      expect(ext.version).not_to be_nil
+      expect(ext.name).to eq "pg_trgm"
+    ensure
+      connection.executed_statements.clear
+      connection.drop_extension(:pg_trgm, if_exists: true)
+      expect(connection.executed_statements).to eq ["DROP EXTENSION IF EXISTS pg_trgm"]
+      expect(connection.extension(:pg_trgm)).to be_nil
+    end
+
+    it "doesn't try to recreate" do
+      connection.create_extension(:pg_trgm, schema: "public")
+      connection.create_extension(:pg_trgm, schema: "public", if_not_exists: true)
+      expect(connection.executed_statements).to eq ["CREATE EXTENSION pg_trgm SCHEMA public",
+                                                    "CREATE EXTENSION IF NOT EXISTS pg_trgm SCHEMA public"]
+    ensure
+      connection.drop_extension(:pg_trgm, if_exists: true)
+    end
+
+    context "non-executing" do
+      around do |example|
+        connection.dont_execute(&example)
+      end
+
+      it "supports all options on create" do
+        connection.create_extension(:my_extension, if_not_exists: true, schema: "public", version: "2.0", cascade: true)
+        expect(connection.executed_statements).to eq(
+          ["CREATE EXTENSION IF NOT EXISTS my_extension SCHEMA public VERSION '2.0' CASCADE"]
+        )
+      end
+
+      it "supports all options on drop" do
+        connection.drop_extension(:my_extension, if_exists: true, cascade: true)
+        expect(connection.executed_statements).to eq ["DROP EXTENSION IF EXISTS my_extension CASCADE"]
+      end
+
+      it "can update an extensions" do
+        connection.alter_extension(:my_extension, version: true)
+        expect(connection.executed_statements).to eq ["ALTER EXTENSION my_extension UPDATE"]
+      end
+
+      it "can update to a specific version" do
+        connection.alter_extension(:my_extension, version: "2.0")
+        expect(connection.executed_statements).to eq ["ALTER EXTENSION my_extension UPDATE TO '2.0'"]
+      end
+
+      it "can change schemas" do
+        connection.alter_extension(:my_extension, schema: "my_app")
+        expect(connection.executed_statements).to eq ["ALTER EXTENSION my_extension SET SCHEMA my_app"]
+      end
+
+      it "cannot change schema and update" do
+        expect { connection.alter_extension(:my_extension, schema: "my_app", version: "2.0") }
+          .to raise_error(ArgumentError)
+      end
+
+      it "can drop multiple extensions" do
+        connection.drop_extension(:my_extension1, :my_extension2)
+        expect(connection.executed_statements).to eq ["DROP EXTENSION my_extension1, my_extension2"]
+      end
+
+      it "does not allow dropping no extensions" do
+        expect { connection.drop_extension }.to raise_error(ArgumentError)
+      end
+    end
+  end
+
+  describe "#add_schema_to_search_path" do
+    around do |example|
+      original_search_path = connection.schema_search_path
+      connection.schema_search_path = "public"
+      example.call
+    ensure
+      connection.schema_search_path = original_search_path
+    end
+
+    it "adds a schema to search path" do
+      connection.add_schema_to_search_path("postgis") do
+        expect(connection.schema_search_path).to eq "public,postgis"
+      end
+      expect(connection.schema_search_path).to eq "public"
+    end
+
+    it "doesn't duplicate an existing schema" do
+      connection.add_schema_to_search_path("public") do
+        expect(connection.schema_search_path).to eq "public"
+      end
+      expect(connection.schema_search_path).to eq "public"
+    end
+
+    it "is cleaned up properly when the transaction rolls back manually" do
+      expect do
+        connection.add_schema_to_search_path("postgis") do
+          raise ActiveRecord::Rollback
+        end
+      end.to raise_error(ActiveRecord::Rollback)
+      expect(connection.schema_search_path).to eq "public"
+    end
+
+    it "is cleaned up properly when the transaction rolls back" do
+      expect do
+        connection.add_schema_to_search_path("postgis") do
+          connection.execute("gibberish")
+        end
+      end.to raise_error(ActiveRecord::StatementInvalid)
+      expect(connection.schema_search_path).to eq "public"
+    end
+  end
+
+  describe "#vacuum" do
+    it "does a straight vacuum of everything" do
+      connection.vacuum
+      expect(connection.executed_statements).to eq ["VACUUM"]
+    end
+
+    it "supports several options" do
+      connection.vacuum(analyze: true, verbose: true)
+      expect(connection.executed_statements).to eq ["VACUUM VERBOSE ANALYZE"]
+    end
+
+    it "validates parallel option is an integer" do
+      expect { connection.vacuum(parallel: :garbage) }.to raise_error(ArgumentError)
+    end
+
+    it "validates parallel option is postive" do
+      expect { connection.vacuum(parallel: -1) }.to raise_error(ArgumentError)
+    end
+
+    context "non-executing" do
+      around do |example|
+        connection.dont_execute(&example)
+      end
+
+      it "vacuums a table" do
+        connection.vacuum(:my_table)
+        expect(connection.executed_statements).to eq ['VACUUM "my_table"']
+      end
+
+      it "vacuums multiple tables" do
+        connection.vacuum(:table1, :table2)
+        expect(connection.executed_statements).to eq ['VACUUM "table1", "table2"']
+      end
+
+      it "requires analyze with specific columns" do
+        expect { connection.vacuum({ my_table: :column1 }) }.to raise_error(ArgumentError)
+      end
+
+      it "analyzes a specific column" do
+        connection.vacuum({ my_table: :column }, analyze: true)
+        expect(connection.executed_statements).to eq ['VACUUM ANALYZE "my_table" ("column")']
+      end
+
+      it "analyzes multiples columns" do
+        connection.vacuum({ my_table: %i[column1 column2] }, analyze: true)
+        expect(connection.executed_statements).to eq ['VACUUM ANALYZE "my_table" ("column1", "column2")']
+      end
+
+      it "analyzes a mixture of tables and columns" do
+        connection.vacuum(:table1, { my_table: %i[column1 column2] }, analyze: true)
+        expect(connection.executed_statements).to eq ['VACUUM ANALYZE "table1", "my_table" ("column1", "column2")']
+      end
+    end
+
+    describe "#wal_lsn_diff" do
+      skip unless connection.wal?
+
+      it "executes" do
+        expect(connection.wal_lsn_diff(:current, :current)).to eq 0
+      end
+    end
+
+    describe "#in_recovery?" do
+      it "works" do
+        expect(connection.in_recovery?).to eq false
+      end
+    end
+  end
 end

data/spec/spec_helper.rb

@@ -3,6 +3,7 @@
 require "activerecord-pg-extensions"
 require "byebug"
 require "active_record/railtie"
+require "active_record/pg_extensions/all"
 
 ActiveRecord::Base # rubocop:disable Lint/Void
 Rails.env = "test"
@@ -26,9 +27,22 @@ module StatementCaptureConnection
     @executed_statements ||= []
   end
 
-  def execute(statement, *)
-    executed_statements << statement
-    super unless @dont_execute
+  %w[execute exec_no_cache exec_cache].each do |method|
+    class_eval <<-RUBY, __FILE__, __LINE__ + 1
+      def #{method}(statement, *)
+        materialize_transactions # this still needs to get called, even if we skip actually executing
+        executed_statements << statement
+        return empty_pg_result if @dont_execute
+
+        super
+      end
+    RUBY
+  end
+
+  # we can't actually generate a dummy one of these, so we just query the db with something
+  # that won't return anything
+  def empty_pg_result
+    @connection.async_exec("SELECT 0 WHERE FALSE")
   end
 end
 ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(StatementCaptureConnection)

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-pg-extensions
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - Cody Cutrer
@@ -160,10 +160,14 @@ files:
 - README.md
 - config/database.yml
 - lib/active_record/pg_extensions.rb
+- lib/active_record/pg_extensions/all.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/version.rb
 - lib/activerecord-pg-extensions.rb
+- spec/pessimistic_migrations_spec.rb
 - spec/postgresql_adapter_spec.rb
 - spec/spec_helper.rb
 homepage: https://github.com/instructure/activerecord-pg-extensions
@@ -193,4 +197,5 @@ summary: Several extensions to ActiveRecord's PostgreSQLAdapter.
 test_files:
 - spec/spec_helper.rb
 - spec/postgresql_adapter_spec.rb
+- spec/pessimistic_migrations_spec.rb
 - config/database.yml