Skip to content

Commit

Permalink
fix: Refactor new connection
Browse files Browse the repository at this point in the history
The old method was calling `AbstractAdapter#initialize`
with a connection object, rather than a hash. This was
then setting `@unconfigured_connection`, which in turns
forced to re-run `#configure_connection` when a first
query was made. This is usually OK (although unnecessary)
except if `use_follower_reads_for_type_introspection` is
used, since it sets AOST for a single query, which fails
within a transaction.

Fixes #320
  • Loading branch information
BuonOmo authored and rafiss committed Apr 22, 2024
1 parent 97ce388 commit 07f9fef
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 27 deletions.
1 change: 1 addition & 0 deletions Gemfile
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ group :development, :test do
gem "debug"
gem "minitest-excludes", "~> 2.0.1"
gem "minitest-github_action_reporter", github: "BuonOmo/minitest-github_action_reporter", require: "minitest/github_action_reporter_plugin"
gem "ostruct", "~> 0.6"

# Gems used for tests meta-programming.
gem "parser"
Expand Down
49 changes: 22 additions & 27 deletions lib/active_record/connection_adapters/cockroachdb_adapter.rb
Original file line number Diff line number Diff line change
Expand Up @@ -32,35 +32,15 @@
ActiveRecord::ConnectionAdapters::CockroachDB.initial_setup

module ActiveRecord
# TODO: once in rails 7.2, remove this and replace with a `#register` call.
# See: https://github.com/rails/rails/commit/22a26d7f74ea8f0d5f7c4169531ae38441cfd5e5#diff-2468c670eb10c24bd2823e42708489a336d6f21c6efc7e3c4a574166fa77bb22
module ConnectionHandling
def cockroachdb_adapter_class
ConnectionAdapters::CockroachDBAdapter
end

def cockroachdb_connection(config)
# This is copied from the PostgreSQL adapter.
conn_params = config.symbolize_keys.compact

# Map ActiveRecords param names to PGs.
conn_params[:user] = conn_params.delete(:username) if conn_params[:username]
conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database]

# Forward only valid config params to PG::Connection.connect.
valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:requiressl]
conn_params.slice!(*valid_conn_param_keys)

ConnectionAdapters::CockroachDBAdapter.new(
ConnectionAdapters::CockroachDBAdapter.new_client(conn_params),
logger,
conn_params,
config
)
# This rescue flow appears in new_client, but it is needed here as well
# since Cockroach will sometimes not raise until a query is made.
rescue ActiveRecord::StatementInvalid => error
no_db_err_check1 = conn_params && conn_params[:dbname] && error.cause.message.include?(conn_params[:dbname])
no_db_err_check2 = conn_params && conn_params[:dbname] && error.cause.message.include?("pg_type")
if no_db_err_check1 || no_db_err_check2
raise ActiveRecord::NoDatabaseError
else
raise ActiveRecord::ConnectionNotEstablished, error.message
end
cockroachdb_adapter_class.new(config)
end
end
end
Expand Down Expand Up @@ -255,6 +235,21 @@ def self.database_exists?(config)
false
end

def initialize(...)
super

# This rescue flow appears in new_client, but it is needed here as well
# since Cockroach will sometimes not raise until a query is made.
rescue ActiveRecord::StatementInvalid => error
no_db_err_check1 = @connection_parameters && @connection_parameters[:dbname] && error.cause.message.include?(@connection_parameters[:dbname])
no_db_err_check2 = @connection_parameters && @connection_parameters[:dbname] && error.cause.message.include?("pg_type")
if no_db_err_check1 || no_db_err_check2
raise ActiveRecord::NoDatabaseError
else
raise ActiveRecord::ConnectionNotEstablished, error.message
end
end

# override
# The PostgreSQLAdapter uses syntax for an anonymous function
# (DO $$) that CockroachDB does not support.
Expand Down

0 comments on commit 07f9fef

Please sign in to comment.