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

1 files changed, 20 insertions, 0 deletions

diff --git a/lib/active_record/database_tasks_extensions.rb b/lib/active_record/database_tasks_extensions.rb
new file mode 100644
index 000000000..e274f476d
--- /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