Skip to content

Commit

Permalink
fix(schema_statements): add unique_rowid to pk_and_sequence_for
Browse files Browse the repository at this point in the history
Allow to get the primary key in the default case.
  • Loading branch information
BuonOmo committed Jan 27, 2024
1 parent 2e06103 commit c6e953e
Show file tree
Hide file tree
Showing 7 changed files with 167 additions and 35 deletions.
2 changes: 1 addition & 1 deletion .ruby-version
Original file line number Diff line number Diff line change
@@ -1 +1 @@
3.2.1
3.2.3
3 changes: 3 additions & 0 deletions Gemfile
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@ group :development, :test do
gem "minitest-excludes", "~> 2.0.1"
gem "minitest-github_action_reporter", github: "BuonOmo/minitest-github_action_reporter", require: "minitest/github_action_reporter_plugin"

# Gems used for tests meta-programming.
gem "parser"

# Gems used by the ActiveRecord test suite
gem "bcrypt", "~> 3.1.18"
gem "mocha", "~> 1.14.0"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,63 @@ def primary_key(table_name)
end
end

# OVERRIDE: Added `unique_rowid` to the last line of the second query.
# This is a CockroachDB-specific function used for primary keys.
#
# Returns a table's primary key and belonging sequence.
def pk_and_sequence_for(table) # :nodoc:
# First try looking for a sequence with a dependency on the
# given table's primary key.
result = query(<<~SQL, "SCHEMA")[0]
SELECT attr.attname, nsp.nspname, seq.relname
FROM pg_class seq,
pg_attribute attr,
pg_depend dep,
pg_constraint cons,
pg_namespace nsp
WHERE seq.oid = dep.objid
AND seq.relkind = 'S'
AND attr.attrelid = dep.refobjid
AND attr.attnum = dep.refobjsubid
AND attr.attrelid = cons.conrelid
AND attr.attnum = cons.conkey[1]
AND seq.relnamespace = nsp.oid
AND cons.contype = 'p'
AND dep.classid = 'pg_class'::regclass
AND dep.refobjid = #{quote(quote_table_name(table))}::regclass
SQL

if result.nil? || result.empty?
result = query(<<~SQL, "SCHEMA")[0]
SELECT attr.attname, nsp.nspname,
CASE
WHEN pg_get_expr(def.adbin, def.adrelid) !~* 'nextval' THEN NULL
WHEN split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2) ~ '.' THEN
substr(split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2),
strpos(split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2), '.')+1)
ELSE split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2)
END
FROM pg_class t
JOIN pg_attribute attr ON (t.oid = attrelid)
JOIN pg_attrdef def ON (adrelid = attrelid AND adnum = attnum)
JOIN pg_constraint cons ON (conrelid = adrelid AND adnum = conkey[1])
JOIN pg_namespace nsp ON (t.relnamespace = nsp.oid)
WHERE t.oid = #{quote(quote_table_name(table))}::regclass
AND cons.contype = 'p'
AND pg_get_expr(def.adbin, def.adrelid) ~* 'nextval|uuid_generate|gen_random_uuid|unique_rowid'
SQL
end

pk = result.shift
if result.last
[pk, PostgreSQL::Name.new(*result)]
else
[pk, nil]
end
rescue
nil
end

# override
# Modified version of the postgresql foreign_keys method.
# Replaces t2.oid::regclass::text with t2.relname since this is
Expand Down
24 changes: 20 additions & 4 deletions test/cases/adapters/postgresql/postgresql_adapter_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
require "cases/helper"
require "support/ddl_helper"
require "support/connection_helper"
require "support/copy_cat"

module CockroachDB
module ConnectionAdapters
Expand All @@ -12,7 +13,6 @@ class PostgreSQLAdapterTest < ActiveRecord::PostgreSQLTestCase

def setup
@connection = ActiveRecord::Base.connection
@connection_handler = ActiveRecord::Base.connection_handler
end

def teardown
Expand Down Expand Up @@ -87,11 +87,27 @@ def test_invalid_index
end
end

# OVERRIDE: the `default_sequence_name` is `nil`, let's `to_s` it
# for a fair comparison.
CopyCat.copy_methods(self, ActiveRecord::ConnectionAdapters::PostgreSQLAdapterTest,
:test_pk_and_sequence_for,
:test_pk_and_sequence_for_with_non_standard_primary_key
) do
attr_accessor :already_updated
def on_send(node)
return super unless node in [:send, _, :default_sequence_name, [:str, "ex"], [:str, "id"|"code"]]

raise "The source code must have changed" if already_updated
already_updated ||= :yes
insert_after(node.loc.expression, ".to_s")
end
end

private

def with_example_table(definition = "id serial primary key, number integer, data character varying(255)", &block)
super(@connection, "ex", definition, &block)
end
CopyCat.copy_methods(self, ActiveRecord::ConnectionAdapters::PostgreSQLAdapterTest,
:with_example_table
)
end
end
end
50 changes: 20 additions & 30 deletions test/cases/migration_test.rb
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
require "cases/helper_cockroachdb"
require "models/person"
require "support/copy_cat"

class Reminder < ActiveRecord::Base; end unless Object.const_defined?(:Reminder)
class Thing < ActiveRecord::Base; end unless Object.const_defined?(:Thing)
Expand Down Expand Up @@ -92,7 +93,6 @@ def migrate(x)

class BulkAlterTableMigrationsTest < ActiveRecord::TestCase
def setup
@original_test = ::BulkAlterTableMigrationsTest.new(name)
@connection = Person.connection
@connection.create_table(:delete_me, force: true) { |t| }
Person.reset_column_information
Expand All @@ -103,38 +103,28 @@ def setup
Person.connection.drop_table(:delete_me) rescue nil
end

%i(
test_adding_indexes
test_removing_index
test_adding_multiple_columns
test_changing_index
).each do |method_name|
file, line = ::BulkAlterTableMigrationsTest.instance_method(method_name).source_location
iter = File.foreach(file)
(line - 1).times { iter.next }
indent = iter.next[/\A\s*/]
content = +""
content << iter.next until iter.peek == indent + "end\n"
content['"PostgreSQLAdapter" =>'] = '"CockroachDBAdapter" =>'
eval(<<-RUBY, binding, __FILE__, __LINE__ + 1)
def #{method_name}
#{content}
end
RUBY
# Change expected query count from PostgreSQLAdapter to CockroachDBAdapter.
CopyCat.copy_methods(self, ::BulkAlterTableMigrationsTest,
:test_adding_indexes,
:test_removing_index,
:test_adding_multiple_columns,
:test_changing_index
) do
def on_str(node)
return unless node in [:str, "PostgreSQLAdapter"]

replace(node.loc.expression, '"CockroachDBAdapter"')
end
end

private

%i[
with_bulk_change_table
column
columns
index
indexes
].each do |method|
define_method(method) do |*args, &block|
@original_test.send(method, *args, &block)
end
end
CopyCat.copy_methods(self, ::BulkAlterTableMigrationsTest,
:with_bulk_change_table,
:column,
:columns,
:index,
:indexes
)
end
end
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
exclude :test_partial_index_on_column_named_like_keyword, "CockroachDB adds parentheses around the WHERE definition."
exclude :test_only_check_for_insensitive_comparison_capability_once, "the DROP DOMAIN syntax is not supported by CRDB"

exclude :test_pk_and_sequence_for, "The test needs a little rework since the sequence is empty in CRDB"
exclude :test_pk_and_sequence_for_with_non_standard_primary_key, "The test needs a little rework since the sequence is empty in CRDB"

plpgsql_needed = "Will be testable with CRDB 23.2, introducing PL-PGSQL"
exclude :test_ignores_warnings_when_behaviour_ignore, plpgsql_needed
Expand Down
64 changes: 64 additions & 0 deletions test/support/copy_cat.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
# frozen_string_literal: true

require "parser/current"

module CopyCat
extend self

# Copy methods from The original rails class to our adapter.
# While copying, we can rewrite the source code of the method using
# ast. Use `debug: true` to lead you true that process.
def copy_methods(new_klass, old_klass, *methods, debug: false, &block)
methods.each do |met|
file, _ = old_klass.instance_method(met).source_location
ast = find_method(Parser::CurrentRuby.parse_file(file), met)
code =
if block_given?
source = ast.location.expression.source
buffer = Parser::Source::Buffer.new(met, source: source)
# We need to recompute the ast to have correct locations.
ast = Parser::CurrentRuby.parse(source)

if debug
puts "=" * 80
puts "Rewriter doc: https://www.rubydoc.info/gems/parser/3.3.0.5/Parser/TreeRewriter"
puts "Pattern matching doc: https://docs.ruby-lang.org/en/master/syntax/pattern_matching_rdoc.html"
puts
puts "Method: #{met}"
puts
puts "Source:"
puts buffer.source
puts
puts "AST:"
pp ast
puts
end
rewriter_class = Class.new(Parser::TreeRewriter, &block)
rewriter_class.new.rewrite(buffer, ast)
else
ast.location.expression.source
end
if debug and block_given?
puts "Rewritten source:"
puts code
puts "=" * 80
end
location = caller_locations(3, 1).first
new_klass.class_eval(code, location.absolute_path, location.lineno)
end
end

def find_method(ast, method_name)
method_name = method_name.to_sym
to_search = [ast]
while !to_search.empty?
node = to_search.shift
next unless node.is_a?(Parser::AST::Node)
if node in [:def, ^method_name, *]
return node
end
to_search += node.children
end
return nil
end
end

0 comments on commit c6e953e

Please sign in to comment.