From 9aaaa96d2ff5e27f065375a2544e86afa31a4e13 Mon Sep 17 00:00:00 2001 From: Claire Date: Fri, 19 Mar 2021 02:43:13 +0100 Subject: [PATCH] Use more robust hook for loading timestamp_id function into database (#15919) --- config/application.rb | 1 + ...70920024819_status_ids_to_timestamp_ids.rb | 4 +- .../database_tasks_extensions.rb | 20 +++++++ lib/tasks/db.rake | 56 ------------------- 4 files changed, 23 insertions(+), 58 deletions(-) create mode 100644 lib/active_record/database_tasks_extensions.rb diff --git a/config/application.rb b/config/application.rb index 0960247b3a..3267fa71be 100644 --- a/config/application.rb +++ b/config/application.rb @@ -28,6 +28,7 @@ require_relative '../lib/webpacker/manifest_extensions' require_relative '../lib/webpacker/helper_extensions' require_relative '../lib/action_dispatch/cookie_jar_extensions' require_relative '../lib/rails/engine_extensions' +require_relative '../lib/active_record/database_tasks_extensions' Dotenv::Railtie.load diff --git a/db/migrate/20170920024819_status_ids_to_timestamp_ids.rb b/db/migrate/20170920024819_status_ids_to_timestamp_ids.rb index c10aa2c4f7..8679f8ece8 100644 --- a/db/migrate/20170920024819_status_ids_to_timestamp_ids.rb +++ b/db/migrate/20170920024819_status_ids_to_timestamp_ids.rb @@ -1,7 +1,7 @@ class StatusIdsToTimestampIds < ActiveRecord::Migration[5.1] def up # Prepare the function we will use to generate IDs. - Rake::Task['db:define_timestamp_id'].execute + Mastodon::Snowflake.define_timestamp_id # Set up the statuses.id column to use our timestamp-based IDs. ActiveRecord::Base.connection.execute(<<~SQL) @@ -11,7 +11,7 @@ class StatusIdsToTimestampIds < ActiveRecord::Migration[5.1] SQL # Make sure we have a sequence to use. - Rake::Task['db:ensure_id_sequences_exist'].execute + Mastodon::Snowflake.ensure_id_sequences_exist end def down diff --git a/lib/active_record/database_tasks_extensions.rb b/lib/active_record/database_tasks_extensions.rb new file mode 100644 index 0000000000..e274f476df --- /dev/null +++ b/lib/active_record/database_tasks_extensions.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +require_relative '../mastodon/snowflake' + +module ActiveRecord + module Tasks + module DatabaseTasks + original_load_schema = instance_method(:load_schema) + + define_method(:load_schema) do |db_config, *args| + ActiveRecord::Base.establish_connection(db_config) + Mastodon::Snowflake.define_timestamp_id + + original_load_schema.bind(self).call(db_config, *args) + + Mastodon::Snowflake.ensure_id_sequences_exist + end + end + end +end diff --git a/lib/tasks/db.rake b/lib/tasks/db.rake index f6c9c7eecf..552a02b3f1 100644 --- a/lib/tasks/db.rake +++ b/lib/tasks/db.rake @@ -1,36 +1,5 @@ # frozen_string_literal: true -require_relative '../mastodon/snowflake' - -def each_schema_load_environment - # If we're in development, also run this for the test environment. - # This is a somewhat hacky way to do this, so here's why: - # 1. We have to define this before we load the schema, or we won't - # have a timestamp_id function when we get to it in the schema. - # 2. db:setup calls db:schema:load_if_ruby, which calls - # db:schema:load, which we define above as having a prerequisite - # of this task. - # 3. db:schema:load ends up running - # ActiveRecord::Tasks::DatabaseTasks.load_schema_current, which - # calls a private method `each_current_configuration`, which - # explicitly also does the loading for the `test` environment - # if the current environment is `development`, so we end up - # needing to do the same, and we can't even use the same method - # to do it. - - if Rails.env.development? - test_conf = ActiveRecord::Base.configurations['test'] - - if test_conf['database']&.present? - ActiveRecord::Base.establish_connection(:test) - yield - ActiveRecord::Base.establish_connection(Rails.env.to_sym) - end - end - - yield -end - namespace :db do namespace :migrate do desc 'Setup the db or migrate depending on state of db' @@ -61,29 +30,4 @@ namespace :db do end Rake::Task['db:migrate'].enhance(['db:post_migration_hook']) - - # Before we load the schema, define the timestamp_id function. - # Idiomatically, we might do this in a migration, but then it - # wouldn't end up in schema.rb, so we'd need to figure out a way to - # get it in before doing db:setup as well. This is simpler, and - # ensures it's always in place. - Rake::Task['db:schema:load'].enhance ['db:define_timestamp_id'] - - # After we load the schema, make sure we have sequences for each - # table using timestamp IDs. - Rake::Task['db:schema:load'].enhance do - Rake::Task['db:ensure_id_sequences_exist'].invoke - end - - task :define_timestamp_id do - each_schema_load_environment do - Mastodon::Snowflake.define_timestamp_id - end - end - - task :ensure_id_sequences_exist do - each_schema_load_environment do - Mastodon::Snowflake.ensure_id_sequences_exist - end - end end