Use more robust hook for loading timestamp_id function into database (#15919)

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
 

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

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

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