如何通用改写Rails默认db:… Rake任务以适配多PostgreSQL Schema？

我正在把一个基于PostgreSQL Schema（靠search_path实现多租户）的Rails应用，迁移到Rails内置的多数据库/水平分片方案。现在需要让rails db:migrate、rails db:migrate:rollback这类所有db相关任务，自动遍历所有租户Schema执行。

之前试过两种方案都有坑：

1. 扩展ActiveRecord::Tasks::DatabaseTasks
   只对db:migrate有用，像db:migrate:down、db:rollback这类任务根本不会调用migrate方法，完全无效：

   # config/initializers/migrations.rb
   module Migrations
     module DatabaseTasks
       def migrate(*)
         ["a", "b", "c"].each do |schema|
           ActiveRecord::Base.connection.schema_search_path = schema
           puts "------------- migrating: #{schema} --------------"
           super
         end
       end
     end
   end
   
   ActiveSupport.on_load(:active_record) do
     ActiveRecord::Tasks::DatabaseTasks.singleton_class.prepend(Migrations::DatabaseTasks)
   end
   

2. 逐个覆盖Rake任务
   要给每个db任务单独写重复执行逻辑，不仅繁琐，任务覆盖的逻辑还容易出问题：

   # tasks/my_custom_db_tasks.rake
   namespace :db do
     task migrate: :environment do |_task, _args|
       original_migrate_task = Rake::Task['db:migrate']
       ["a", "b", "c"].each do |schema|
         ActiveRecord::Base.connection.schema_search_path = schema
         puts "----------  migrating: #{schema} --------------"
         original_migrate_task.invoke
         original_migrate_task.reenable
       end
     end
   end
   

现在想找个通用的方式，给所有db:...任务挂钩子，让它们自动在指定的所有Schema上重复执行原逻辑。

通用解决方案：动态包装所有Rake任务 #

可以写一个通用的Rake任务包装器，自动遍历所有Schema并执行原任务，不用逐个去覆盖。

1. 基础版本：自动包装所有db任务 #

在lib/tasks/multi_schema.rake里添加以下代码：

namespace :multi_schema do
  def run_for_all_schemas(schemas = ["a", "b", "c"])
    # 获取当前要执行的原始任务名称
    original_task_name = Rake.application.top_level_tasks.first
    original_task = Rake::Task[original_task_name]

    schemas.each do |schema|
      puts "=== 执行 #{original_task_name} 针对 Schema: #{schema} ==="
      # 切换到目标Schema
      ActiveRecord::Base.connection.schema_search_path = schema
      # 执行原始任务
      original_task.execute
      # 重新启用任务，确保能多次执行
      original_task.reenable
    end
  end

  # 遍历所有db任务，动态创建包装任务
  Rake::Task.tasks.each do |task|
    # 跳过自身命名空间的任务，避免循环
    next unless task.name.start_with?("db:") && !task.name.include?("multi_schema")

    task_parts = task.name.split(":")
    # 拆出命名空间和任务名
    namespace_name = task_parts[0..-2].join(":")
    task_name = task_parts.last

    # 重新定义同名任务，替换原任务
    namespace namespace_name.to_sym do
      task task_name.to_sym => :environment do
        run_for_all_schemas
      end
    end
  end
end

2. 优化版本：支持传递命令行参数 #

上面的基础版本处理不了带参数的任务（比如db:migrate:down VERSION=20240520123456），可以优化参数传递逻辑：

namespace :multi_schema do
  def run_for_all_schemas(schemas = ["a", "b", "c"])
    original_task_name = Rake.application.top_level_tasks.first
    original_task = Rake::Task[original_task_name]
    # 提取命令行中的参数（比如VERSION=xxx）
    task_args = ARGV.select { |arg| arg.include?("=") }.join(" ")

    schemas.each do |schema|
      puts "=== 执行 #{original_task_name} #{task_args} 针对 Schema: #{schema} ==="
      ActiveRecord::Base.connection.schema_search_path = schema
      # 带参数执行任务
      original_task.execute(ARGV)
      original_task.reenable
    end
  end

  # 动态包装所有db任务
  Rake::Task.tasks.each do |task|
    next unless task.name.start_with?("db:") && !task.name.include?("multi_schema")

    task_parts = task.name.split(":")
    current_namespace = self
    # 逐层创建命名空间
    task_parts[0..-2].each do |ns|
      current_namespace = current_namespace.namespace(ns.to_sym) unless current_namespace.const_defined?(ns.camelize)
    end

    # 定义包装后的任务
    current_namespace.send(:task, task_parts.last.to_sym => :environment) do
      run_for_all_schemas
    end
  end
end

3. 适配Rails内置分片的方案 #

如果已经在迁移到Rails的多数据库分片，可以直接利用connected_to切换分片，不用手动改schema_search_path：

namespace :multi_shard do
  def run_for_all_shards
    # 假设你的分片配置在database.yml里是tenant_a、tenant_b这类命名
    shard_names = Rails.application.config.database_configuration.keys.select { |k| k.start_with?("tenant_") }
    
    original_task_name = Rake.application.top_level_tasks.first
    original_task = Rake::Task[original_task_name]
    task_args = ARGV.select { |arg| arg.include?("=") }

    shard_names.each do |shard|
      puts "=== 执行 #{original_task_name} 针对分片: #{shard} ==="
      # 切换到目标分片
      ActiveRecord::Base.connected_to(role: :writing, shard: shard) do
        original_task.execute(task_args)
        original_task.reenable
      end
    end
  end

  # 同样的方式动态包装所有db任务
  Rake::Task.tasks.each do |task|
    next unless task.name.start_with?("db:") && !task.name.include?("multi_shard")

    task_parts = task.name.split(":")
    current_namespace = self
    task_parts[0..-2].each do |ns|
      current_namespace = current_namespace.namespace(ns.to_sym) unless current_namespace.const_defined?(ns.camelize)
    end

    current_namespace.send(:task, task_parts.last.to_sym => :environment) do
      run_for_all_shards
    end
  end
end

注意事项 #

• 确保这个Rake任务是最后加载的：在Rakefile的末尾添加load 'lib/tasks/multi_schema.rake'，这样才能覆盖原有的db任务。
• 对于db:create这类初始化任务，需要先确保数据库连接已建立，或者单独处理（因为切换Schema前需要数据库存在）。
• 如果Schema/分片列表是动态的（比如从数据库读取租户），可以把固定的数组替换成查询逻辑，但要等environment任务加载完成后再执行查询。

内容的提问来源于stack exchange，提问作者Niels Kristian