From 1c48fa2dc10ed822df3825c1b18c00465ee67418 Mon Sep 17 00:00:00 2001 From: Mehul Mathur Date: Thu, 12 Sep 2024 20:51:51 +0000 Subject: [PATCH] chore: added statement timeout gem --- Gemfile | 1 + Gemfile.lock | 3 + config/initializers/statement_timeout.rb | 3 - lib/statement_timeout.rb | 53 -------- spec/lib/statement_timeout_spec.rb | 163 ----------------------- 5 files changed, 4 insertions(+), 219 deletions(-) delete mode 100644 config/initializers/statement_timeout.rb delete mode 100644 lib/statement_timeout.rb delete mode 100644 spec/lib/statement_timeout_spec.rb diff --git a/Gemfile b/Gemfile index 15fc18d0f..ff5e5c8ec 100644 --- a/Gemfile +++ b/Gemfile @@ -21,6 +21,7 @@ gem 'redis', '~> 4.7.1' # API migrations gem 'request_migrations', '~> 1.1' gem 'verbose_migrations' +gem 'statement_timeout' # API params gem 'typed_params', '~> 1.2.5' diff --git a/Gemfile.lock b/Gemfile.lock index e86da498a..80cca93b2 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -473,6 +473,8 @@ GEM activesupport (>= 5.2) sprockets (>= 3.0.0) stackprof (0.2.26) + statement_timeout (1.0.0) + rails (>= 6.0) statsd-ruby (1.5.0) stringio (3.1.0) stripe (5.48.0) @@ -573,6 +575,7 @@ DEPENDENCIES sidekiq-cronitor (~> 3.6.0) sprockets (~> 3.0) stackprof + statement_timeout stripe (~> 5.43) stripe-ruby-mock! strong_migrations diff --git a/config/initializers/statement_timeout.rb b/config/initializers/statement_timeout.rb deleted file mode 100644 index 0ebaf6ebf..000000000 --- a/config/initializers/statement_timeout.rb +++ /dev/null @@ -1,3 +0,0 @@ -# frozen_string_literal: true - -require_dependency Rails.root / 'lib' / 'statement_timeout' diff --git a/lib/statement_timeout.rb b/lib/statement_timeout.rb deleted file mode 100644 index 5f7f4dd79..000000000 --- a/lib/statement_timeout.rb +++ /dev/null @@ -1,53 +0,0 @@ -# frozen_string_literal: true - -module StatementTimeout - module AbstractAdapterExtension - def supports_statement_timeout? = false - def statement_timeout = raise NotImplementedError - def statement_timeout=(timeout) - raise NotImplementedError - end - end - - module PostgreSQLAdapterExtension - def supports_statement_timeout? = true - def statement_timeout = @statement_timeout ||= query_value("SHOW statement_timeout") - def statement_timeout=(timeout) - @statement_timeout = nil - - internal_exec_query("SET statement_timeout = #{quote(timeout)}") - end - end - - module QueryMethodsExtension - def statement_timeout(timeout) - timeout = if timeout in ActiveSupport::Duration - timeout.in_milliseconds - else - timeout - end - - connection_pool.with_connection do |connection| - raise ActiveRecord::AdapterError, "statement_timeout is not supported for the #{connection.class.inspect} adapter" unless - connection.supports_statement_timeout? - - statement_timeout_was, connection.statement_timeout = connection.statement_timeout, timeout - - yield connection - ensure - connection.statement_timeout = statement_timeout_was - end - end - end - - module QueryingExtension - delegate :statement_timeout, to: :all - end - - ActiveSupport.on_load :active_record do - ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(AbstractAdapterExtension) - ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(PostgreSQLAdapterExtension) - ActiveRecord::QueryMethods.prepend(QueryMethodsExtension) - ActiveRecord::Querying.prepend(QueryingExtension) - end -end diff --git a/spec/lib/statement_timeout_spec.rb b/spec/lib/statement_timeout_spec.rb deleted file mode 100644 index 6ad91108c..000000000 --- a/spec/lib/statement_timeout_spec.rb +++ /dev/null @@ -1,163 +0,0 @@ -# frozen_string_literal: true - -require 'rails_helper' -require 'spec_helper' - -require_dependency Rails.root / 'lib' / 'statement_timeout' - -describe StatementTimeout do - before(:all) { DatabaseCleaner.strategy = nil } # vacuum doesn't support transactions - after(:all) { DatabaseCleaner.strategy = :transaction } - - subject { License } - - let(:connection) { subject.respond_to?(:lease_connection) ? subject.lease_connection : subject.connection } - let(:table_name) { subject.table_name } - - context 'with a duration' do - it 'should set a temporary statement_timeout for valid duration' do - statement_timeout_was = connection.statement_timeout - - expect { subject.statement_timeout(1.minute) { subject.unscoped.take } }.to( - match_queries(count: 3) do |queries| - expect(queries.first).to eq <<~SQL.squish - SET statement_timeout = 60000 - SQL - - expect(queries.second).to eq <<~SQL.squish - SELECT "#{table_name}".* FROM "#{table_name}" LIMIT 1 - SQL - - expect(queries.third).to eq <<~SQL.squish - SET statement_timeout = '#{statement_timeout_was}' - SQL - end - ) - - expect(connection.statement_timeout).to eq statement_timeout_was - end - end - - context 'with an integer' do - it 'should set a temporary statement_timeout for valid integer' do - statement_timeout_was = connection.statement_timeout - - expect { subject.statement_timeout(1000) { subject.unscoped.take } }.to( - match_queries(count: 3) do |queries| - expect(queries.first).to eq <<~SQL.squish - SET statement_timeout = 1000 - SQL - - expect(queries.second).to eq <<~SQL.squish - SELECT "#{table_name}".* FROM "#{table_name}" LIMIT 1 - SQL - - expect(queries.third).to eq <<~SQL.squish - SET statement_timeout = '#{statement_timeout_was}' - SQL - end - ) - - expect(connection.statement_timeout).to eq statement_timeout_was - end - end - - context 'with a float' do - it 'should set a temporary statement_timeout for valid float' do - statement_timeout_was = connection.statement_timeout - - expect { subject.statement_timeout(1000.5) { subject.unscoped.take } }.to( - match_queries(count: 3) do |queries| - expect(queries.first).to eq <<~SQL.squish - SET statement_timeout = 1000.5 - SQL - - expect(queries.second).to eq <<~SQL.squish - SELECT "#{table_name}".* FROM "#{table_name}" LIMIT 1 - SQL - - expect(queries.third).to eq <<~SQL.squish - SET statement_timeout = '#{statement_timeout_was}' - SQL - end - ) - - expect(connection.statement_timeout).to eq statement_timeout_was - end - end - - context 'with a string' do - it 'should set a temporary statement_timeout for valid value' do - statement_timeout_was = connection.statement_timeout - - expect { subject.statement_timeout('1s') { subject.unscoped.take } }.to( - match_queries(count: 3) do |queries| - expect(queries.first).to eq <<~SQL.squish - SET statement_timeout = '1s' - SQL - - expect(queries.second).to eq <<~SQL.squish - SELECT "#{table_name}".* FROM "#{table_name}" LIMIT 1 - SQL - - expect(queries.third).to eq <<~SQL.squish - SET statement_timeout = '#{statement_timeout_was}' - SQL - end - ) - - expect(connection.statement_timeout).to eq statement_timeout_was - end - - it 'should raise for invalid value' do - statement_timeout_was = connection.statement_timeout - - expect { subject.statement_timeout('foo') { subject.unscoped.take } } - .to raise_error ActiveRecord::StatementInvalid - - expect(connection.statement_timeout).to eq statement_timeout_was - end - end - - it 'should not timeout' do - expect { subject.statement_timeout('2s') { connection.execute('select pg_sleep(1)') } } - .to_not raise_error - end - - it 'should timeout' do - expect { subject.statement_timeout('1s') { connection.execute('select pg_sleep(2)') } } - .to raise_error ActiveRecord::StatementInvalid - end - - it 'should return a relation' do - expect(subject.statement_timeout('1s') { subject.all }).to be_an ActiveRecord::Relation - end - - it 'should return a record' do - expect(subject.statement_timeout('1s') { subject.new }).to be_a subject - end - - it 'should return a value' do - expect(subject.statement_timeout('1s') { connection.execute('select 1 as value')[0]['value'] }).to eq 1 - end - - it 'should support transaction' do - expect { subject.statement_timeout('1s') { subject.transaction { subject.unscoped.take } } } - .to_not raise_error - end - - # NOTE(ezekg) We're explicitly testing VACUUM because it doesn't support - # transactions, and this asserts our implementation is - # compatible with such queries. - it 'should support vacuum' do - expect { subject.statement_timeout('1s') { connection.execute("VACUUM ANALYZE #{table_name}") } } - .to_not raise_error - end - - it 'should raise error' do - expect { subject.statement_timeout('1s') { subject.transaction { connection.execute("VACUUM ANALYZE #{table_name}") } } } - .to raise_error { |err| - expect(err.message).to match /vacuum cannot run inside a transaction block/i - } - end -end