Rails Postgresql Query timeout set and skip during migration

Hi, I am Shubham Mishra from India,  this is the part of Ruby on Rails exploration journey.  In this post we will discuss about the rails DB (Postgresql) query time out setting and how can we avoid this timeout setting during migration and I am sure you will get excited to make your hands dirty with those code, so why to wait let's get started...

Question:

Rails and postgresql I want to set query timeout to 10 sec but skip to check this timeout setting during rails migration how can i achieve?

Answer:

To set a query timeout of 10 seconds for PostgreSQL in a Rails application, you can use the statement_timeout configuration option in your database.yml file. Here is an example configuration that sets the query timeout to 10 seconds:

	production:
	adapter: postgresql
	database: my_app_production
	username: my_app
	password: secret
	host: localhost
	port: 5432
	variables:
	statement_timeout: 10000 # Set query timeout to 10 seconds

view raw rails_pg_config.yml hosted with ❤ by GitHub

This will set the query timeout to 10 seconds for all queries executed by the PostgreSQL adapter in the production environment.

To skip checking the query timeout setting during Rails migrations, you can use a custom connection adapter that sets the statement_timeout to nil during migrations. Here is an example implementation of a custom adapter that does this:

	# config/initializers/postgresql_adapter.rb
	module ActiveRecord
	module ConnectionAdapters
	class PostgreSQLAdapter < AbstractAdapter
	def exec_query(sql, name = nil, binds = [])
	if migration_context? # Skip query timeout check during migrations
	super(sql, name, binds)
	else
	with_statement_timeout(10_000) { super(sql, name, binds) } # Set query timeout to 10 seconds
	end
	end
	private
	def with_statement_timeout(timeout)
	previous_timeout = @connection.exec("SHOW statement_timeout").first["statement_timeout"]
	@connection.exec("SET statement_timeout = #{timeout}")
	yield
	ensure
	@connection.exec("SET statement_timeout = #{previous_timeout}")
	end
	def migration_context?
	Thread.current[:__migration_context__] || false
	end
	end
	end
	end

view raw rails_postgres_customer_adapter.rb hosted with ❤ by GitHub

This custom adapter overrides the exec_query method of the PostgreSQLAdapter class to set the statement_timeout to 10 seconds for all queries executed outside of migrations, and to skip the timeout check during migrations. The with_statement_timeout method is used to set the timeout temporarily and ensure that it is reset to the previous value after the query is executed.

To use this custom adapter during migrations, you can set the __migration_context__ thread-local variable to true before running the migration. For example:

	# db/migrate/20230419000000_create_products.rb
	Thread.current[:__migration_context__] = true
	class CreateProducts < ActiveRecord::Migration[6.1]
	def change
	create_table :products do |t|
	t.string :name
	t.integer :price
	t.timestamps
	end
	end
	end
	Thread.current[:__migration_context__] = false

view raw postgres_migration_with_flag.rb hosted with ❤ by GitHub

This will ensure that the custom adapter skips the query timeout check during the migration, and sets the timeout to 10 seconds for all other queries executed by the PostgreSQL adapter.

Another Approach: Now in case you dont want to go for these many changes, u can use a database connection that has a different configuration specifically for migrations. This way, you can set the query timeout to a higher value or disable it altogether for migrations, without affecting the configuration for the rest of the application. Here's an example configuration that sets a query timeout of 10 seconds for the production environment, but disables the query timeout for migrations:

	production:
	adapter: postgresql
	database: my_app_production
	username: my_app
	password: secret
	host: localhost
	port: 5432
	variables:
	statement_timeout: 10000 # Set query timeout to 10 seconds
	production_for_migrations:
	adapter: postgresql
	database: my_app_production
	username: my_app
	password: secret
	host: localhost
	port: 5432

view raw diffrent_pg_config_rails_for_migration.yml hosted with ❤ by GitHub

In this configuration, the production_for_migrations connection uses the same PostgreSQL database as production, but does not set the statement_timeout variable, effectively disabling the query timeout.

To use the production_for_migrations connection for migrations, you can set the DATABASE_URL environment variable to the URL of this connection when running migrations. For example:

$ DATABASE_URL=postgres://my_app:secret@localhost/my_app_production_for_migrations rails db:migrate

OR

In the deployment script you can mention Rails Env with the db migrate cmd which will make sure to run the migration with mentioned confg connection,

RAILS_ENV=production_for_migrations rake db:migrate

This will ensure that the query timeout is disabled for migrations, while still using the normal configuration with the query timeout set to 10 seconds for the rest of the application.

But wait, I have one more solution, is you can have seprate config for normal DB queries and migrations timeout in rakefile, rakefile it always be there on the root path of your rails application. And the final code would look like this:

	# Add your own tasks in files placed in lib/tasks ending in .rake,
	# for example lib/tasks/capistrano.rake, and they will automatically be available to Rake.
	require File.expand_path('../config/application', __FILE__)
	namespace :db do
	task :setup do
	# Set statement timeout for each query
	if ActiveRecord::Base.connection_config[:adapter] == 'postgresql'
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';") # 10 seconds
	end
	end
	task :migrate => :environment do
	# Disable statement timeout during migration
	if ActiveRecord::Base.connection_config[:adapter] == 'postgresql'
	ActiveRecord::Base.connection.execute("SET statement_timeout = '0';")
	end
	ActiveRecord::MigrationContext.new(['db/migrate'], ActiveRecord::SchemaMigration).migrate
	end
	end
	Rails.application.load_tasks

view raw rakefile_with_db_timeout_without_migration.rb hosted with ❤ by GitHub

You can change the value 10000 to any other value and 0 timeout for migration block is saying no timeout.

One more solution:

In this solution what you need to do is to create a initializer file and register callbacks for migrations, before_migration and after_migration,

	if ActiveRecord::Base.connection_config[:adapter] == 'postgresql'
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';") # 10 seconds
	# Register an before migration callback to reset the timeout to 0 which means no timeout
	ActiveRecord::Migration.before_migrate do
	ActiveRecord::Base.connection.execute("SET statement_timeout = '0';")
	end
	ActiveRecord::Migration.after_migrate do
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';")
	end
	end

view raw before_after_migration_callbacks.rb hosted with ❤ by GitHub

I think this is best solution rather than touching rakefile

Actually in my case i had Rails 7.0.4.3 so above solution doesn't worked as those callbacks were not available for my version so u can try below solutions:

	# config/initializers/active_record_query_timeout.rb
	# Set statement timeout for each query
	if ActiveRecord::Base.connection_config[:adapter] == 'postgresql'
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';") # 10 seconds
	ActiveRecord::MigrationContext.prepend(Module.new do
	def initialize(*args)
	super
	ActiveRecord::Base.connection.execute("SET statement_timeout = '0';")
	end
	def iterate_with_statement_timeout(*args, &block)
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';")
	super
	end
	end)
	ActiveRecord::MigrationProxy.prepend(Module.new do
	def run_migration(*args)
	migration = migration
	ActiveRecord::MigrationContext.new(migration.version, migration.name, nil).iterate_with_statement_timeout do
	super
	end
	end
	end)
	end
	#One more soulution you can try if above doesn't support for your verion:
	# config/initializers/active_record_query_timeout.rb
	if ActiveRecord::Base.connection_config[:adapter] == 'postgresql'
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';") # 10 seconds
	# Register a before migration callback to reset the timeout to 0 which means no timeout
	ActiveRecord::MigrationContext.before_each_migration do |migration|
	ActiveRecord::Base.connection.execute("SET statement_timeout = '0';")
	end
	# Register an after migration callback to reset the timeout
	ActiveRecord::MigrationContext.after_each_migration do |migration|
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';")
	end
	end
	#OR
	# config/initializers/active_record_query_timeout.rb
	# Set statement timeout for each query
	if ActiveRecord::Base.connection_config[:adapter] == 'postgresql'
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';") # 10 seconds
	# Define a method to reset the timeout to 10 seconds
	def reset_query_timeout
	ActiveRecord::Base.connection.execute("SET statement_timeout = '10000';")
	end
	# Add a before_action to set the timeout to 0 before the migration
	ActiveSupport.on_load(:active_record) do
	ActiveRecord::Migration.before_each_migration do |migration|
	ActiveRecord::Base.connection.execute("SET statement_timeout = '0';")
	end
	ActiveRecord::Migration.after_each_migration do |migration|
	reset_query_timeout
	end
	end
	end

view raw timeout_DB_query_except_migration_rails_7.0.4.3.rb hosted with ❤ by GitHub

You can keep on trying by changing the migration timeout to '1' to get time out for rake db:migrate

You can explore my previous blog: DB query optimization in Ruby on Rails

Or you can explore my other blogs too here,

Get to know answers for common search on Google : A blog for posts which can help you for daily life problems, such as where to get free images, Topic suggestion for the blog.

Computer Science algorithms and other knowledge share : A blog for posts such as best search algorithm, Top interview questions for diffrent technologies, knowledge share for some frameworks or programming languages for the interview or in general terms.

My ideas to solve real world problems : A blog where me shared and presented my ideas to solve a real world problems, this will be interesting for me.

Ruby on Rails Web development Blog : As the name suggest, it is the blog for sharing few knowledge about RoR web development framework.

Liked my blogs, wanna to connect:

LinkedIn  |  GitHub  |  HackerEarth  |  ResearchGate  |  Twitter  |  Facebook  |  StackOverflow

Thanks for reading this post, Have a good day :)