How can I set statement_timeout for an individual migration?

Question

I'd like to set postgres statement_timeout for an individual migration. I can't seem to be able to do this. Here's my experiment:

def change
  execute <<~SQL
    SET LOCAL statement_timeout = 1; -- ms

    -- this does not cause a timeout which is expected, because pg
    -- only applies the timeout to the next protocol message / statement,
    -- and rails sends everthing inside execute in the same statement
    select pg_sleep(1); -- seconds
  SQL

  # if uncommented, this DOES cause a timeout, which is expected
  # execute <<~SQL
  #   select pg_sleep(1); -- seconds
  # SQL

  # this does not cause a timeout, which is unexpected
  remove_column :foos, :bar

  # we do get here, which is unexpected
  raise "we finished"
end

How can I do this?

Answer 1

I assume your question is about setting statement_timeout for an individual statement within your migration (not for an individual migration).

You can achieve this by using Postgres’s SET instead of SET LOCAL.

class SomeMigration < ActiveRecord::Migration[6.0]
  # Disables BEGIN/COMMIT around migration. SET LOCAL will have no effect.
  disable_ddl_transaction!

  def change
    # Applies to the current database session, until changed or reset.
    execute <<~SQL
      SET statement_timeout = '1s';
    SQL

    # Will raise `ActiveRecord::QueryCanceled`.
    execute <<~SQL
      SELECT pg_sleep(2);
    SQL

    # Resets (disables) timeout.
    execute <<~SQL
      SET statement_timeout = DEFAULT;
    SQL

    # Will not raise.
    execute <<~SQL
      SELECT pg_sleep(2);
    SQL
  end
end

Note that you could use SET LOCAL if you don’t call disable_ddl_transaction!. Combining both doesn’t work because SET LOCAL outside a transaction block has no effect. It logs:

WARNING: SET LOCAL can only be used in transaction blocks

PS: Your call to disable_ddl_transaction! would raise NoMethodError. It should be before the def change definition.