chore: various backports to 7.2

.github/workflows/ci.yml

@@ -41,7 +41,7 @@ jobs:
       matrix:
         # https://www.cockroachlabs.com/docs/releases/release-support-policy
         crdb: [v23.2, v24.1, v24.2]
-        ruby: ["3.3"]
+        ruby: [3.4]
     name: Test (crdb=${{ matrix.crdb }} ruby=${{ matrix.ruby }})
     steps:
       - name: Set Up Actions
@@ -88,4 +88,7 @@ jobs:
           done
           cat ${{ github.workspace }}/setup.sql | cockroach sql --insecure
       - name: Test
-        run: bundle exec rake test TESTOPTS='--profile=5'
+        run: bundle exec rake test
+        env:
+          TESTOPTS: "--profile=5"
+          RAILS_MINITEST_PLUGIN: "1" # Make sure that we use the minitest plugin for profiling.

CHANGELOG.md

@@ -1,6 +1,8 @@
 # Changelog
 
-## Ongoing
+## 7.2.1 - 2025-03-26
+
+- Fix transaction state on rollback ([#364](https://github.com/cockroachdb/activerecord-cockroachdb-adapter/pull/364))
 
 ## 7.2.0 - 2024-09-24
 

CONTRIBUTING.md

@@ -117,7 +117,11 @@ This section intent to help you with a checklist.
 - Check for some important methods, some will change for sure:
   - [x] `def new_column_from_field(`
   - [x] `def column_definitions(`
-  - [x] `def pk_and_sequence_for(`
+  - [x] # `def pk_and_sequence_for(`
+  - [ ] `def new_column_from_field(`
+  - [ ] `def column_definitions(`
+  - [ ] `def pk_and_sequence_for(`
+  - [ ] `def foreign_keys(` and `def all_foreign_keys(`
   - [ ] ...
 - Check for setups containing `drop_table` in the test suite.
   Especially if you have tons of failure, this is likely the cause.

bin/console

@@ -12,20 +12,22 @@ require "activerecord-cockroachdb-adapter"
 #       structure_load(Post.connection_db_config, "awesome-file.sql")
 require "active_record/connection_adapters/cockroachdb/database_tasks"
 
+DB_NAME = "ar_crdb_console"
+
 schema_kind = ENV.fetch("SCHEMA_KIND", ENV.fetch("SCHEMA", "default"))
 
-system("cockroach sql --insecure --host=localhost:26257 --execute='drop database if exists ar_crdb_console'",
+system("cockroach sql --insecure --host=localhost:26257 --execute='drop database if exists #{DB_NAME}'",
   exception: true)
-system("cockroach sql --insecure --host=localhost:26257 --execute='create database ar_crdb_console'",
+system("cockroach sql --insecure --host=localhost:26257 --execute='create database #{DB_NAME}'",
   exception: true)
 
 ActiveRecord::Base.establish_connection(
-  #Alternative version:  "cockroachdb://root@localhost:26257/ar_crdb_console"
+  #Alternative version:  "cockroachdb://root@localhost:26257/#{DB_NAME}"
   adapter: "cockroachdb",
   host: "localhost",
   port: 26257,
   user: "root",
-  database: "ar_crdb_console"
+  database: DB_NAME
 )
 
 load "#{__dir__}/console_schemas/#{schema_kind}.rb"

bin/console_schemas/default.rb

@@ -6,4 +6,6 @@ class Post < ActiveRecord::Base
     t.string :title
     t.text :body
   end
+
+  add_index("posts", ["title"], name: "index_posts_on_title", unique: true)
 end

lib/active_record/connection_adapters/cockroachdb/database_statements.rb

@@ -24,9 +24,17 @@ def insert_fixtures_set(fixture_set, tables_to_delete = [])
           table_deletes = tables_to_delete.map { |table| "DELETE FROM #{quote_table_name(table)}" }
           statements = table_deletes + fixture_inserts
 
-          with_multi_statements do
-            disable_referential_integrity do
-              execute_batch(statements, "Fixtures Load")
+          begin # much faster without disabling referential integrity, worth trying.
+            with_multi_statements do
+              transaction(requires_new: true) do
+                execute_batch(statements, "Fixtures Load")
+              end
+            end
+          rescue
+            with_multi_statements do
+              disable_referential_integrity do
+                execute_batch(statements, "Fixtures Load")
+              end
             end
           end
         end

lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb

@@ -34,11 +34,18 @@ def check_all_foreign_keys_valid!
         end
 
         def disable_referential_integrity
-          foreign_keys = tables.map { |table| foreign_keys(table) }.flatten
+          foreign_keys = all_foreign_keys
 
-          foreign_keys.each do |foreign_key|
-            remove_foreign_key(foreign_key.from_table, name: foreign_key.options[:name])
+          statements = foreign_keys.map do |foreign_key|
+            # We do not use the `#remove_foreign_key` method here because it
+            # checks for foreign keys existance in the schema cache. This method
+            # is performance critical and we know the foreign key exist.
+            at = create_alter_table foreign_key.from_table
+            at.drop_foreign_key foreign_key.name
+
+            schema_creation.accept(at)
           end
+          execute_batch(statements, "Disable referential integrity -> remove foreign keys")
 
           yield
 
@@ -52,19 +59,88 @@ def disable_referential_integrity
           ActiveRecord::Base.table_name_suffix = ""
 
           begin
-            foreign_keys.each do |foreign_key|
-              # Avoid having PG:DuplicateObject error if a test is ran in transaction.
-              # TODO: verify that there is no cache issue related to running this (e.g: fk
-              #   still in cache but not in db)
-              next if foreign_key_exists?(foreign_key.from_table, name: foreign_key.options[:name])
+            # Avoid having PG:DuplicateObject error if a test is ran in transaction.
+            # TODO: verify that there is no cache issue related to running this (e.g: fk
+            #   still in cache but not in db)
+            #
+            # We avoid using `foreign_key_exists?` here because it checks the schema cache
+            # for every key. This method is performance critical for the test suite, hence
+            # we use the `#all_foreign_keys` method that only make one query to the database.
+            already_inserted_foreign_keys = all_foreign_keys
+            statements = foreign_keys.map do |foreign_key|
+              next if already_inserted_foreign_keys.any? { |fk| fk.from_table == foreign_key.from_table && fk.options[:name] == foreign_key.options[:name] }
+
+              options = foreign_key_options(foreign_key.from_table, foreign_key.to_table, foreign_key.options)
+              at = create_alter_table foreign_key.from_table
+              at.add_foreign_key foreign_key.to_table, options
 
-              add_foreign_key(foreign_key.from_table, foreign_key.to_table, **foreign_key.options)
+              schema_creation.accept(at)
             end
+            execute_batch(statements.compact, "Disable referential integrity -> add foreign keys")
           ensure
             ActiveRecord::Base.table_name_prefix = old_prefix
             ActiveRecord::Base.table_name_suffix = old_suffix
           end
         end
+
+        private
+
+        # Copy/paste of the `#foreign_keys(table)` method adapted to return every single
+        # foreign key in the database.
+        def all_foreign_keys
+          fk_info = exec_query(<<~SQL, "SCHEMA")
+            SELECT CASE
+              WHEN n1.nspname = current_schema()
+              THEN ''
+              ELSE n1.nspname || '.'
+            END || t1.relname AS from_table,
+            CASE
+              WHEN n2.nspname = current_schema()
+              THEN ''
+              ELSE n2.nspname || '.'
+            END || t2.relname AS to_table,
+            a1.attname AS column, a2.attname AS primary_key, c.conname AS name, c.confupdtype AS on_update, c.confdeltype AS on_delete, c.convalidated AS valid, c.condeferrable AS deferrable, c.condeferred AS deferred,
+            c.conkey, c.confkey, c.conrelid, c.confrelid
+            FROM pg_constraint c
+            JOIN pg_class t1 ON c.conrelid = t1.oid
+            JOIN pg_class t2 ON c.confrelid = t2.oid
+            JOIN pg_attribute a1 ON a1.attnum = c.conkey[1] AND a1.attrelid = t1.oid
+            JOIN pg_attribute a2 ON a2.attnum = c.confkey[1] AND a2.attrelid = t2.oid
+            JOIN pg_namespace t3 ON c.connamespace = t3.oid
+            JOIN pg_namespace n1 ON t1.relnamespace = n1.oid
+            JOIN pg_namespace n2 ON t2.relnamespace = n2.oid
+            WHERE c.contype = 'f'
+            ORDER BY c.conname
+          SQL
+
+          fk_info.map do |row|
+            from_table = PostgreSQL::Utils.unquote_identifier(row["from_table"])
+            to_table = PostgreSQL::Utils.unquote_identifier(row["to_table"])
+            conkey = row["conkey"].scan(/\d+/).map(&:to_i)
+            confkey = row["confkey"].scan(/\d+/).map(&:to_i)
+
+            if conkey.size > 1
+              column = column_names_from_column_numbers(row["conrelid"], conkey)
+              primary_key = column_names_from_column_numbers(row["confrelid"], confkey)
+            else
+              column = PostgreSQL::Utils.unquote_identifier(row["column"])
+              primary_key = row["primary_key"]
+            end
+
+            options = {
+              column: column,
+              name: row["name"],
+              primary_key: primary_key
+            }
+            options[:on_delete] = extract_foreign_key_action(row["on_delete"])
+            options[:on_update] = extract_foreign_key_action(row["on_update"])
+            options[:deferrable] = extract_constraint_deferrable(row["deferrable"], row["deferred"])
+
+            options[:validate] = row["valid"]
+
+            ForeignKeyDefinition.new(from_table, to_table, options)
+          end
+        end
       end
     end
   end

lib/active_record/connection_adapters/cockroachdb/schema_statements.rb

@@ -55,6 +55,37 @@ def primary_key(table_name)
           end
         end
 
+        def primary_keys(table_name)
+          return super unless database_version >= 24_02_02
+
+          query_values(<<~SQL, "SCHEMA")
+            SELECT a.attname
+              FROM (
+                     SELECT indrelid, indkey, generate_subscripts(indkey, 1) idx
+                       FROM pg_index
+                      WHERE indrelid = #{quote(quote_table_name(table_name))}::regclass
+                        AND indisprimary
+                   ) i
+              JOIN pg_attribute a
+                ON a.attrelid = i.indrelid
+               AND a.attnum = i.indkey[i.idx]
+               AND NOT a.attishidden
+             ORDER BY i.idx
+          SQL
+        end
+
+        def column_names_from_column_numbers(table_oid, column_numbers)
+          return super unless database_version >= 24_02_02
+
+          Hash[query(<<~SQL, "SCHEMA")].values_at(*column_numbers).compact
+            SELECT a.attnum, a.attname
+            FROM pg_attribute a
+            WHERE a.attrelid = #{table_oid}
+            AND a.attnum IN (#{column_numbers.join(", ")})
+            AND NOT a.attishidden
+          SQL
+        end
+
         # OVERRIDE: CockroachDB does not support deferrable constraints.
         #   See: https://go.crdb.dev/issue-v/31632/v23.1
         def foreign_key_options(from_table, to_table, options)
@@ -265,31 +296,6 @@ def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, **) #
           sql
         end
 
-        # This overrides the method from PostegreSQL adapter
-        # Resets the sequence of a table's primary key to the maximum value.
-        def reset_pk_sequence!(table, pk = nil, sequence = nil)
-          unless pk && sequence
-            default_pk, default_sequence = pk_and_sequence_for(table)
-
-            pk ||= default_pk
-            sequence ||= default_sequence
-          end
-
-          if @logger && pk && !sequence
-            @logger.warn "#{table} has primary key #{pk} with no default sequence."
-          end
-
-          if pk && sequence
-            quoted_sequence = quote_table_name(sequence)
-            max_pk = query_value("SELECT MAX(#{quote_column_name pk}) FROM #{quote_table_name(table)}", "SCHEMA")
-            if max_pk.nil?
-              minvalue = query_value("SELECT seqmin FROM pg_sequence WHERE seqrelid = #{quote(quoted_sequence)}::regclass", "SCHEMA")
-            end
-
-            query_value("SELECT setval(#{quote(quoted_sequence)}, #{max_pk ? max_pk : minvalue}, #{max_pk ? true : false})", "SCHEMA")
-          end
-        end
-
         # override
         def native_database_types
           # Add spatial types

lib/active_record/connection_adapters/cockroachdb/transaction_manager.rb

@@ -45,6 +45,27 @@ def within_new_transaction(isolation: nil, joinable: true, attempts: 0)
           within_new_transaction(isolation: isolation, joinable: joinable, attempts: attempts + 1) { yield }
         end
 
+        # OVERRIDE: the `rescue ActiveRecord::StatementInvalid` block is new, see comment.
+        def rollback_transaction(transaction = nil)
+          @connection.lock.synchronize do
+            transaction ||= @stack.last
+            begin
+              transaction.rollback
+            rescue ActiveRecord::StatementInvalid => err
+              # This is important to make Active Record aware the record was not inserted/saved
+              # Otherwise Active Record will assume save was successful and it doesn't retry the transaction
+              # See this thread for more details:
+              # https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/258#issuecomment-2256633329
+              transaction.rollback_records if err.cause.is_a?(PG::NoActiveSqlTransaction)
+
+              raise
+            ensure
+              @stack.pop if @stack.last == transaction
+            end
+            transaction.rollback_records
+          end
+        end
+
         def retryable?(error)
           return true if serialization_error?(error)
           return true if error.is_a? ActiveRecord::SerializationFailure

lib/version.rb

@@ -15,5 +15,5 @@
 # limitations under the License.
 
 module ActiveRecord
-  COCKROACH_DB_ADAPTER_VERSION = "7.2.0"
+  COCKROACH_DB_ADAPTER_VERSION = "7.2.1"
 end

test/cases/migration/hidden_column_test.rb

@@ -53,6 +53,17 @@ def test_add_hidden_column
           output = dump_table_schema "rockets"
           assert_match %r{t.uuid "new_col", hidden: true$}, output
         end
+
+        # Since 24.2.2, hash sharded indexes add a hidden column to the table.
+        # This tests ensure that the user can still drop the index even if they
+        # call `#remove_index` with the column name rather than the index name.
+        def test_remove_index_with_a_hidden_column
+          @connection.execute <<-SQL
+            CREATE INDEX hash_idx ON rockets (name) USING HASH WITH (bucket_count=8);
+          SQL
+          @connection.remove_index :rockets, :name
+          assert :ok
+        end
       end
     end
   end

test/cases/primary_keys_test.rb

@@ -97,4 +97,38 @@ def test_schema_dump_primary_key_integer_with_default_nil
       assert_match %r{create_table "int_defaults", id: :bigint, default: nil}, schema
     end
   end
+
+  class PrimaryKeyHiddenColumnTest < ActiveRecord::TestCase
+    class Rocket < ActiveRecord::Base
+    end
+
+    def setup
+      connection = ActiveRecord::Base.connection
+      connection.execute <<-SQL
+        CREATE TABLE rockets(
+          id SERIAL PRIMARY KEY USING HASH WITH (bucket_count=4)
+        )
+      SQL
+    end
+
+    def teardown
+      ActiveRecord::Base.connection.drop_table :rockets
+    end
+
+    def test_to_key_with_hidden_primary_key_part
+      rocket = Rocket.new
+      assert_nil rocket.to_key
+      rocket.save
+      assert_equal rocket.to_key, [rocket.id]
+    end
+
+    def test_read_attribute_with_hidden_primary_key_part
+      rocket = Rocket.create!
+      id = assert_not_deprecated(ActiveRecord.deprecator) do
+        rocket.read_attribute(:id)
+      end
+
+      assert_equal rocket.id, id
+    end
+  end
 end