From 7d9c15d7469229cdd8a8997d309ba98b26710aa3 Mon Sep 17 00:00:00 2001 From: Adrianna Chang Date: Fri, 15 Mar 2024 16:39:12 -0400 Subject: [PATCH] Retry known idempotent SELECT queries on connection-related exceptions This commit makes two types of queries retry-able by opting into our `allow_retry` flag: 1) SELECT queries we construct by walking the Arel tree via `#to_sql_and_binds`. We use a new `retryable` attribute on collector classes, which defaults to true for most node types, but will be set to false for non-idempotent node types (functions, SQL literals, etc). The `retryable` value is returned from `#to_sql_and_binds` and used by `#select_all` and passed down the call stack, eventually reaching the adapter's `#internal_exec_query` method. 2) `#find` and `#find_by` queries with known attributes. We set `allow_retry: true` in `#cached_find_by`, and pass this down to `#find_by_sql` and `#_query_by_sql`. These changes ensure that queries we know are safe to retry can be retried automatically. --- activerecord/CHANGELOG.md | 9 +++ .../abstract/database_statements.rb | 24 ++++--- .../abstract/query_cache.rb | 8 +-- .../abstract_mysql_adapter.rb | 4 +- .../mysql2/database_statements.rb | 4 +- .../connection_adapters/postgresql_adapter.rb | 4 +- .../sqlite3/database_statements.rb | 2 +- .../trilogy/database_statements.rb | 4 +- activerecord/lib/active_record/core.rb | 2 +- activerecord/lib/active_record/querying.rb | 8 +-- .../lib/active_record/statement_cache.rb | 6 +- activerecord/lib/arel/collectors/composite.rb | 2 +- .../lib/arel/collectors/sql_string.rb | 2 +- .../lib/arel/collectors/substitute_binds.rb | 2 +- activerecord/lib/arel/visitors/to_sql.rb | 6 ++ activerecord/test/cases/adapter_test.rb | 66 ++++++++++++++++++- .../test/cases/arel/visitors/to_sql_test.rb | 48 ++++++++++++++ 17 files changed, 166 insertions(+), 35 deletions(-) diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md index 6dcb56afc2033..7c782ffc7eda4 100644 --- a/activerecord/CHANGELOG.md +++ b/activerecord/CHANGELOG.md @@ -1,3 +1,12 @@ +* Retry known idempotent SELECT queries on connection-related exceptions + + SELECT queries we construct by walking the Arel tree and / or with known model attributes + are idempotent and can safely be retried in the case of a connection error. Previously, + adapters such as `TrilogyAdapter` would raise `ActiveRecord::ConnectionFailed: Trilogy::EOFError` + when encountering a connection error mid-request. + + *Adrianna Chang* + * Add dirties option to uncached This adds a `dirties` option to `ActiveRecord::Base.uncached` and diff --git a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb index 62e925a31bbd3..d8244ff61e178 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb @@ -14,7 +14,7 @@ def to_sql(arel_or_sql_string, binds = []) sql end - def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc: + def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil, allow_retry = false) # :nodoc: # Arel::TreeManager -> Arel::Node if arel_or_sql_string.respond_to?(:ast) arel_or_sql_string = arel_or_sql_string.ast @@ -27,6 +27,7 @@ def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc: end collector = collector() + collector.retryable = true if prepared_statements collector.preparable = true @@ -41,10 +42,11 @@ def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc: else sql = visitor.compile(arel_or_sql_string, collector) end - [sql.freeze, binds, preparable] + allow_retry = collector.retryable + [sql.freeze, binds, preparable, allow_retry] else arel_or_sql_string = arel_or_sql_string.dup.freeze unless arel_or_sql_string.frozen? - [arel_or_sql_string, binds, preparable] + [arel_or_sql_string, binds, preparable, allow_retry] end end private :to_sql_and_binds @@ -64,11 +66,15 @@ def cacheable_query(klass, arel) # :nodoc: end # Returns an ActiveRecord::Result instance. - def select_all(arel, name = nil, binds = [], preparable: nil, async: false) + def select_all(arel, name = nil, binds = [], preparable: nil, async: false, allow_retry: false) arel = arel_from_relation(arel) - sql, binds, preparable = to_sql_and_binds(arel, binds, preparable) + sql, binds, preparable, allow_retry = to_sql_and_binds(arel, binds, preparable, allow_retry) - select(sql, name, binds, prepare: prepared_statements && preparable, async: async && FutureResult::SelectAll) + select(sql, name, binds, + prepare: prepared_statements && preparable, + async: async && FutureResult::SelectAll, + allow_retry: allow_retry + ) rescue ::RangeError ActiveRecord::Result.empty(async: async) end @@ -507,7 +513,7 @@ def high_precision_current_timestamp HIGH_PRECISION_CURRENT_TIMESTAMP end - def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc: + def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc: raise NotImplementedError end @@ -606,7 +612,7 @@ def combine_multi_statements(total_sql) end # Returns an ActiveRecord::Result instance. - def select(sql, name = nil, binds = [], prepare: false, async: false) + def select(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false) if async && async_enabled? if current_transaction.joinable? raise AsynchronousQueryInsideTransactionError, "Asynchronous queries are not allowed inside transactions" @@ -627,7 +633,7 @@ def select(sql, name = nil, binds = [], prepare: false, async: false) return future_result end - result = internal_exec_query(sql, name, binds, prepare: prepare) + result = internal_exec_query(sql, name, binds, prepare: prepare, allow_retry: allow_retry) if async FutureResult.wrap(result) else diff --git a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb index c52fdf32a74c7..9edae461af043 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb @@ -204,19 +204,19 @@ def clear_query_cache pool.clear_query_cache end - def select_all(arel, name = nil, binds = [], preparable: nil, async: false) # :nodoc: + def select_all(arel, name = nil, binds = [], preparable: nil, async: false, allow_retry: false) # :nodoc: arel = arel_from_relation(arel) # If arel is locked this is a SELECT ... FOR UPDATE or somesuch. # Such queries should not be cached. if @query_cache&.enabled? && !(arel.respond_to?(:locked) && arel.locked) - sql, binds, preparable = to_sql_and_binds(arel, binds, preparable) + sql, binds, preparable, allow_retry = to_sql_and_binds(arel, binds, preparable) if async - result = lookup_sql_cache(sql, name, binds) || super(sql, name, binds, preparable: preparable, async: async) + result = lookup_sql_cache(sql, name, binds) || super(sql, name, binds, preparable: preparable, async: async, allow_retry: allow_retry) FutureResult.wrap(result) else - cache_sql(sql, name, binds) { super(sql, name, binds, preparable: preparable, async: async) } + cache_sql(sql, name, binds) { super(sql, name, binds, preparable: preparable, async: async, allow_retry: allow_retry) } end else super diff --git a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb index 3791e0f3a1842..5c8dfb867e655 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb @@ -229,12 +229,12 @@ def disable_referential_integrity # :nodoc: # Mysql2Adapter doesn't have to free a result after using it, but we use this method # to write stuff in an abstract way without concerning ourselves about whether it # needs to be explicitly freed or not. - def execute_and_free(sql, name = nil, async: false) # :nodoc: + def execute_and_free(sql, name = nil, async: false, allow_retry: false) # :nodoc: sql = transform_query(sql) check_if_write_query(sql) mark_transaction_written_if_write(sql) - yield raw_execute(sql, name, async: async) + yield raw_execute(sql, name, async: async, allow_retry: allow_retry) end def begin_db_transaction # :nodoc: diff --git a/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb b/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb index ce79cd5a88531..7d4e4105da2ef 100644 --- a/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb +++ b/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb @@ -18,9 +18,9 @@ def select_all(*, **) # :nodoc: result end - def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc: + def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc: if without_prepared_statement?(binds) - execute_and_free(sql, name, async: async) do |result| + execute_and_free(sql, name, async: async, allow_retry: allow_retry) do |result| if result build_result(columns: result.fields, rows: result.to_a) else diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb index 8dcfc85767f90..f1962ff494095 100644 --- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb +++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb @@ -881,7 +881,7 @@ def exec_no_cache(sql, name, binds, async:, allow_retry:, materialize_transactio type_casted_binds = type_casted_binds(binds) log(sql, name, binds, type_casted_binds, async: async) do |notification_payload| - with_raw_connection(allow_retry: false, materialize_transactions: materialize_transactions) do |conn| + with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn| result = conn.exec_params(sql, type_casted_binds) verified! notification_payload[:row_count] = result.count @@ -895,7 +895,7 @@ def exec_cache(sql, name, binds, async:, allow_retry:, materialize_transactions: update_typemap_for_default_timezone - with_raw_connection(allow_retry: false, materialize_transactions: materialize_transactions) do |conn| + with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn| stmt_key = prepare_statement(sql, binds, conn) type_casted_binds = type_casted_binds(binds) diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb b/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb index a91f08b6f15f7..feffa32690448 100644 --- a/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb +++ b/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb @@ -21,7 +21,7 @@ def explain(arel, binds = [], _options = []) SQLite3::ExplainPrettyPrinter.new.pp(result) end - def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: false) # :nodoc: + def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false) # :nodoc: sql = transform_query(sql) check_if_write_query(sql) diff --git a/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb b/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb index 8f145b155e6a9..9349dc038c25b 100644 --- a/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb +++ b/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb @@ -12,12 +12,12 @@ def select_all(*, **) # :nodoc: result end - def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc: + def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc: sql = transform_query(sql) check_if_write_query(sql) mark_transaction_written_if_write(sql) - result = raw_execute(sql, name, async: async) + result = raw_execute(sql, name, async: async, allow_retry: allow_retry) ActiveRecord::Result.new(result.fields, result.to_a) end diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb index dad3354ef2bb6..7510cea6aa9c2 100644 --- a/activerecord/lib/active_record/core.rb +++ b/activerecord/lib/active_record/core.rb @@ -431,7 +431,7 @@ def cached_find_by(keys, values) } begin - statement.execute(values.flatten, lease_connection).first + statement.execute(values.flatten, lease_connection, allow_retry: true).first rescue TypeError raise ActiveRecord::StatementInvalid end diff --git a/activerecord/lib/active_record/querying.rb b/activerecord/lib/active_record/querying.rb index 18a3f92c8e13a..27d0603f3a9d9 100644 --- a/activerecord/lib/active_record/querying.rb +++ b/activerecord/lib/active_record/querying.rb @@ -47,8 +47,8 @@ module Querying # # Note that building your own SQL query string from user input may expose your application to # injection attacks (https://guides.rubyonrails.org/security.html#sql-injection). - def find_by_sql(sql, binds = [], preparable: nil, &block) - _load_from_sql(_query_by_sql(sql, binds, preparable: preparable), &block) + def find_by_sql(sql, binds = [], preparable: nil, allow_retry: false, &block) + _load_from_sql(_query_by_sql(sql, binds, preparable: preparable, allow_retry: allow_retry), &block) end # Same as #find_by_sql but perform the query asynchronously and returns an ActiveRecord::Promise. @@ -58,8 +58,8 @@ def async_find_by_sql(sql, binds = [], preparable: nil, &block) end end - def _query_by_sql(sql, binds = [], preparable: nil, async: false) # :nodoc: - lease_connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async) + def _query_by_sql(sql, binds = [], preparable: nil, async: false, allow_retry: false) # :nodoc: + lease_connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async, allow_retry: allow_retry) end def _load_from_sql(result_set, &block) # :nodoc: diff --git a/activerecord/lib/active_record/statement_cache.rb b/activerecord/lib/active_record/statement_cache.rb index 6c7359b77b8dc..411a073a72c96 100644 --- a/activerecord/lib/active_record/statement_cache.rb +++ b/activerecord/lib/active_record/statement_cache.rb @@ -62,7 +62,7 @@ def sql_for(binds, connection) end class PartialQueryCollector - attr_accessor :preparable + attr_accessor :preparable, :retryable def initialize @parts = [] @@ -142,12 +142,12 @@ def initialize(query_builder, bind_map, klass) @klass = klass end - def execute(params, connection, &block) + def execute(params, connection, allow_retry: false, &block) bind_values = bind_map.bind params sql = query_builder.sql_for bind_values, connection - klass.find_by_sql(sql, bind_values, preparable: true, &block) + klass.find_by_sql(sql, bind_values, preparable: true, allow_retry: allow_retry, &block) rescue ::RangeError [] end diff --git a/activerecord/lib/arel/collectors/composite.rb b/activerecord/lib/arel/collectors/composite.rb index 0f05dfbe548c3..c8d8d4ccd1a92 100644 --- a/activerecord/lib/arel/collectors/composite.rb +++ b/activerecord/lib/arel/collectors/composite.rb @@ -3,7 +3,7 @@ module Arel # :nodoc: all module Collectors class Composite - attr_accessor :preparable + attr_accessor :preparable, :retryable def initialize(left, right) @left = left diff --git a/activerecord/lib/arel/collectors/sql_string.rb b/activerecord/lib/arel/collectors/sql_string.rb index 8aa8958a1f858..b27eab4f6f2c1 100644 --- a/activerecord/lib/arel/collectors/sql_string.rb +++ b/activerecord/lib/arel/collectors/sql_string.rb @@ -5,7 +5,7 @@ module Arel # :nodoc: all module Collectors class SQLString < PlainString - attr_accessor :preparable + attr_accessor :preparable, :retryable def initialize(*) super diff --git a/activerecord/lib/arel/collectors/substitute_binds.rb b/activerecord/lib/arel/collectors/substitute_binds.rb index 82315c75d321b..020a53054a651 100644 --- a/activerecord/lib/arel/collectors/substitute_binds.rb +++ b/activerecord/lib/arel/collectors/substitute_binds.rb @@ -3,7 +3,7 @@ module Arel # :nodoc: all module Collectors class SubstituteBinds - attr_accessor :preparable + attr_accessor :preparable, :retryable def initialize(quoter, delegate_collector) @quoter = quoter diff --git a/activerecord/lib/arel/visitors/to_sql.rb b/activerecord/lib/arel/visitors/to_sql.rb index ff82d680c268e..46ce3407c907a 100644 --- a/activerecord/lib/arel/visitors/to_sql.rb +++ b/activerecord/lib/arel/visitors/to_sql.rb @@ -20,6 +20,7 @@ def compile(node, collector = Arel::Collectors::SQLString.new) private def visit_Arel_Nodes_DeleteStatement(o, collector) + collector.retryable = false o = prepare_delete_statement(o) if has_join_sources?(o) @@ -37,6 +38,7 @@ def visit_Arel_Nodes_DeleteStatement(o, collector) end def visit_Arel_Nodes_UpdateStatement(o, collector) + collector.retryable = false o = prepare_update_statement(o) collector << "UPDATE " @@ -49,6 +51,7 @@ def visit_Arel_Nodes_UpdateStatement(o, collector) end def visit_Arel_Nodes_InsertStatement(o, collector) + collector.retryable = false collector << "INSERT INTO " collector = visit o.relation, collector @@ -381,6 +384,7 @@ def visit_Arel_Nodes_Group(o, collector) end def visit_Arel_Nodes_NamedFunction(o, collector) + collector.retryable = false collector << o.name collector << "(" collector << "DISTINCT " if o.distinct @@ -768,10 +772,12 @@ def visit_Arel_Nodes_BindParam(o, collector) def visit_Arel_Nodes_SqlLiteral(o, collector) collector.preparable = false + collector.retryable = false collector << o.to_s end def visit_Arel_Nodes_BoundSqlLiteral(o, collector) + collector.retryable = false bind_index = 0 new_bind = lambda do |value| diff --git a/activerecord/test/cases/adapter_test.rb b/activerecord/test/cases/adapter_test.rb index f6e6297dab2a2..75c949c467c7f 100644 --- a/activerecord/test/cases/adapter_test.rb +++ b/activerecord/test/cases/adapter_test.rb @@ -630,18 +630,78 @@ def teardown assert_predicate @connection, :active? end - test "querying after a failed query restores and succeeds" do + test "querying after a failed non-retryable query restores and succeeds" do Post.first # Connection verified (and prepared statement pool populated if enabled) remote_disconnect @connection assert_raises(ActiveRecord::ConnectionFailed) do - Post.first # Connection no longer verified after failed query + @connection.execute("INSERT INTO posts(title, body) VALUES ('foo', 'bar')") end assert Post.first # Verifying the connection causes a reconnect and the query succeeds + assert_predicate @connection, :active? + end + + test "idempotent SELECT queries are retried and result in a reconnect" do + Post.first + remote_disconnect @connection + + assert Post.first assert_predicate @connection, :active? + + remote_disconnect @connection + + assert Post.where(id: [1, 2]).first + assert_predicate @connection, :active? + end + + test "#find and #find_by queries with known attributes are retried and result in a reconnect" do + Post.first + + remote_disconnect @connection + + assert Post.find(1) + assert_predicate @connection, :active? + + remote_disconnect @connection + + assert Post.find_by(title: "Welcome to the weblog") + assert_predicate @connection, :active? + end + + test "queries containing SQL fragments are not retried" do + Post.first + + remote_disconnect @connection + + assert_raises(ActiveRecord::ConnectionFailed) { Post.where("1 = 1").to_a } + assert_not_predicate @connection, :active? + + remote_disconnect @connection + + assert_raises(ActiveRecord::ConnectionFailed) { Post.select("title AS custom_title").first } + assert_not_predicate @connection, :active? + + remote_disconnect @connection + + assert_raises(ActiveRecord::ConnectionFailed) { Post.find_by("updated_at < ?", 2.weeks.ago) } + assert_not_predicate @connection, :active? + end + + test "queries containing SQL functions are not retried" do + Post.first + + remote_disconnect @connection + + tags_count_attr = Post.arel_table[:tags_count] + abs_tags_count = Arel::Nodes::NamedFunction.new("ABS", [tags_count_attr]) + + assert_raises(ActiveRecord::ConnectionFailed) do + Post.where(abs_tags_count.eq(2)).first + end + assert_not_predicate @connection, :active? end test "transaction restores after remote disconnection" do @@ -779,6 +839,8 @@ def raw_transaction_open?(connection) def remote_disconnect(connection) case connection.adapter_name when "PostgreSQL" + # Connection was left in a bad state, need to reconnect to simulate fresh disconnect + connection.verify! if connection.instance_variable_get(:@raw_connection).status == ::PG::CONNECTION_BAD unless connection.instance_variable_get(:@raw_connection).transaction_status == ::PG::PQTRANS_INTRANS connection.instance_variable_get(:@raw_connection).async_exec("begin") end diff --git a/activerecord/test/cases/arel/visitors/to_sql_test.rb b/activerecord/test/cases/arel/visitors/to_sql_test.rb index 7f01ac50575fa..a09c4798d64c5 100644 --- a/activerecord/test/cases/arel/visitors/to_sql_test.rb +++ b/activerecord/test/cases/arel/visitors/to_sql_test.rb @@ -69,6 +69,54 @@ def dispatch _(sql).must_be_like %{ omg(*) IS NULL } end + it "should mark collector as non-retryable when visiting named function" do + function = Nodes::NamedFunction.new("omg", [Arel.star]) + collector = Collectors::SQLString.new + @visitor.accept(function, collector) + + assert_not collector.retryable + end + + it "should mark collector as non-retryable when visiting SQL literal" do + node = Nodes::SqlLiteral.new("COUNT(*)") + collector = Collectors::SQLString.new + @visitor.accept(node, collector) + + assert_not collector.retryable + end + + it "should mark collector as non-retryable when visiting bound SQL literal" do + node = Nodes::BoundSqlLiteral.new("id IN (?)", [[1, 2, 3]], {}) + collector = Collectors::SQLString.new + @visitor.accept(node, collector) + + assert_not collector.retryable + end + + it "should mark collector as non-retryable when visiting insert statement node" do + statement = Arel::Nodes::InsertStatement.new(@table) + collector = Collectors::SQLString.new + @visitor.accept(statement, collector) + + assert_not collector.retryable + end + + it "should mark collector as non-retryable when visiting update statement node" do + statement = Arel::Nodes::UpdateStatement.new(@table) + collector = Collectors::SQLString.new + @visitor.accept(statement, collector) + + assert_not collector.retryable + end + + it "should mark collector as non-retryable when visiting delete statement node" do + statement = Arel::Nodes::DeleteStatement.new(@table) + collector = Collectors::SQLString.new + @visitor.accept(statement, collector) + + assert_not collector.retryable + end + it "should visit built-in functions" do function = Nodes::Count.new([Arel.star]) assert_equal "COUNT(*)", compile(function)