Add tenant shard migration system with registry and migrator service

etewiah · etewiah · commit ee285eabcaa8 · 2026-01-07T17:15:12.000Z
diff --git a/app/lib/pwb/shard_registry.rb b/app/lib/pwb/shard_registry.rb
@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module Pwb
+  module ShardRegistry
+    LOGICAL_TO_PHYSICAL = {
+      default: :primary,
+      shard_1: :tenant_shard_1,
+      shard_2: :tenant_shard_2,
+      demo: :demo_shard
+    }.freeze
+
+    module_function
+
+    def logical_shards
+      LOGICAL_TO_PHYSICAL.keys
+    end
+
+    def physical_name(logical)
+      LOGICAL_TO_PHYSICAL[logical.to_sym]
+    end
+
+    def configured?(logical)
+      physical = physical_name(logical)
+      return false unless physical
+
+      configs.any? { |config| config.name.to_sym == physical }
+    end
+
+    def configs
+      ActiveRecord::Base.configurations.configs_for(env_name: Rails.env)
+    end
+
+    def describe_shard(logical)
+      config = configs.find { |c| c.name.to_sym == physical_name(logical) }
+      return { name: logical, configured: false } unless config
+
+      { name: logical, configured: true, database: config.database, host: config.host }
+    end
+  end
+end
diff --git a/app/services/pwb/tenant_shard_migrator.rb b/app/services/pwb/tenant_shard_migrator.rb
@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+module Pwb
+  class TenantShardMigrator
+    class MigrationError < StandardError; end
+
+    BATCH_SIZE = ENV.fetch('PWB_TENANT_MIGRATION_BATCH', 500).to_i
+
+    attr_reader :website, :target_shard, :logger
+
+    def initialize(website:, target_shard:, logger: Rails.logger, dry_run: false)
+      @website = website
+      @target_shard = normalize_shard(target_shard)
+      @logger = logger
+      @dry_run = dry_run
+    end
+
+    def call
+      validate_target!
+      return if @dry_run
+
+      ActsAsTenant.without_tenant do
+        website.with_lock do
+          migrate_records!
+          website.update!(shard_name: target_shard.to_s)
+        end
+      end
+    end
+
+    private
+
+    def migrate_records!
+      source = website.database_shard
+      logger.info("[TenantShardMigrator] Moving website ##{website.id} from #{source} → #{target_shard}")
+
+      tenant_table_names.each do |table_name|
+        migrated = migrate_table(table_name, source: source, target: target_shard)
+        logger.info("[TenantShardMigrator] #{table_name}: migrated #{migrated} rows") if migrated.positive?
+      end
+    end
+
+    def migrate_table(table_name, source:, target:)
+      total = 0
+
+      loop do
+        rows = fetch_rows(table_name, source)
+        break if rows.empty?
+
+        ids = rows.map { |row| row['id'] }
+        ensure_no_conflicts!(table_name, ids, target: target)
+        insert_rows(table_name, rows, target)
+        delete_rows(table_name, ids, source)
+        total += rows.size
+      end
+
+      total
+    end
+
+    def fetch_rows(table_name, shard)
+      with_connection(shard) do |connection|
+        connection.select_all(<<~SQL, 'TenantShardMigrator').to_a
+          SELECT *
+          FROM #{connection.quote_table_name(table_name)}
+          WHERE website_id = #{connection.quote(website.id)}
+          ORDER BY id ASC
+          LIMIT #{BATCH_SIZE}
+        SQL
+      end
+    end
+
+    def insert_rows(table_name, rows, shard)
+      return if rows.empty?
+
+      with_connection(shard) do |connection|
+        connection.insert_all(rows, table_name)
+      end
+    end
+
+    def delete_rows(table_name, ids, shard)
+      return if ids.empty?
+
+      with_connection(shard) do |connection|
+        sql = <<~SQL.squish
+          DELETE FROM #{connection.quote_table_name(table_name)}
+          WHERE id IN (#{ids.map { |id| connection.quote(id) }.join(', ')})
+        SQL
+        connection.execute(sql)
+      end
+    end
+
+    def ensure_no_conflicts!(table_name, ids, target:)
+      return if ids.empty?
+
+      with_connection(target) do |connection|
+        sql = <<~SQL.squish
+          SELECT 1 FROM #{connection.quote_table_name(table_name)}
+          WHERE id IN (#{ids.map { |id| connection.quote(id) }.join(', ')})
+          LIMIT 1
+        SQL
+        conflict = connection.select_value(sql)
+        raise MigrationError, "ID conflict detected for #{table_name}" if conflict
+      end
+    end
+
+    def with_connection(shard)
+      PwbTenant::ApplicationRecord.connected_to(role: :writing, shard: shard) do
+        yield PwbTenant::ApplicationRecord.connection
+      end
+    end
+
+    def tenant_table_names
+      @tenant_table_names ||= begin
+        base_connection = ActiveRecord::Base.connection
+        base_connection.tables
+                       .reject { |table| table.in?(%w[ar_internal_metadata schema_migrations active_storage_blobs active_storage_attachments]) }
+                       .select do |table|
+                         base_connection.columns(table).any? { |column| column.name == 'website_id' }
+                       end
+      end
+    end
+
+    def validate_target!
+      raise MigrationError, 'Website is already on the target shard' if website.database_shard == target_shard
+      raise MigrationError, "Shard #{target_shard} is not configured" unless physical_shard_configured?(target_shard)
+    end
+
+    def physical_shard_configured?(logical_shard)
+      Pwb::ShardRegistry.configured?(logical_shard)
+    end
+
+    def normalize_shard(value)
+      value = value.to_sym if value.respond_to?(:to_sym)
+      value || :default
+    end
+  end
+end
diff --git a/docs/multi_tenancy/PREMIUM_ENTERPRISE_SHARDING_PLAN.md b/docs/multi_tenancy/PREMIUM_ENTERPRISE_SHARDING_PLAN.md
@@ -117,13 +117,17 @@ Assigns a **new, empty** tenant to a specific shard. Useful for provisioning new
 ### `rake pwb:sharding:migrate[tenant_id, target_shard]`
 **DANGER**: Moves an existing tenant's data from their current shard to the target shard.
 *   **Usage**: `rake pwb:sharding:migrate[123, shard_1]`
+*   **Implementation Status**: ✅ Implemented via `Pwb::TenantShardMigrator`. The task copies every `PwbTenant::` model in batches, verifies there are no ID collisions on the destination, inserts the rows, and then deletes them from the source before updating `website.shard_name`.
+*   **Prerequisites**:
+    * Target shard must be configured in `database.yml` and up to date schema-wise.
+    * Destination shard should not contain colliding primary keys for the migrating tenant tables. (The migrator aborts if conflicts are detected.)
 *   **Steps**:
-    1.  Puts tenant in "Maintenance Mode" (locks site).
-    2.  Dumps data.
-    3.  Restores to target.
-    4.  Verifies row counts match.
-    5.  Updates `website.shard_name`.
-    6.  Unlocks site.
+    1.  Acquire a lock on the website record to prevent concurrent shard changes.
+    2.  Stream tenant data from the source shard in batches.
+    3.  Insert batches into the target shard (raising if any ID conflict is detected).
+    4.  Delete migrated rows from the source shard.
+    5.  Update `website.shard_name`.
+    6.  Release the lock.
 
 ## Tenant Admin UI
 We will add a "Platform Operations" section to the Super Admin interface (`/admin`).
diff --git a/lib/tasks/sharding.rake b/lib/tasks/sharding.rake
@@ -1,39 +1,58 @@
 namespace :pwb do
   namespace :sharding do
-    desc "List all configured shards and tenant counts"
+    desc "List shard configuration, databases, and tenant counts"
     task list: :environment do
-      shards = PwbTenant::ApplicationRecord.connection_handler.connection_pool_names.map { |n| n.split(":").last }
-      # Filter for shards defined in our config if possible, or just used the logical names
-      # Since Rails doesn't easily expose the raw config list at runtime without internal APIs,
-      # we will check the ones we know about or just query websites
-      
-      puts "Configured Shards (from database.yml):"
-      # Rough heuristic to find shard keys
-      configs = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env)
-      configs.each do |config|
-        puts "- #{config.name} (DB: #{config.database})" 
-      end
+      puts format("%-10s %-25s %-12s", "Shard", "Database", "Tenants")
+      puts '-' * 55
 
-      puts "\nTenant Distribution:"
-      Pwb::Website.group(:shard_name).count.each do |shard, count|
-        puts "  #{shard || 'default'}: #{count} tenants"
+      Pwb::ShardRegistry.logical_shards.each do |logical|
+        details = Pwb::ShardRegistry.describe_shard(logical)
+        tenant_count = Pwb::Website.where(shard_name: logical.to_s).count
+        database = details[:configured] ? details[:database] : 'not configured'
+        puts format("%-10s %-25s %-12d", logical, database, tenant_count)
       end
     end
 
     desc "Provision a new tenant on a specific shard"
     task :provision, [:website_id, :shard_name] => :environment do |_, args|
       website = Pwb::Website.find(args[:website_id])
-      target_shard = args[:shard_name]
+      target_shard = args[:shard_name].presence&.to_sym
+
+      unless target_shard
+        puts 'Usage: rake pwb:sharding:provision[website_id,shard_name]'
+        next
+      end
+
+      unless Pwb::ShardRegistry.configured?(target_shard)
+        puts "ERROR: Shard '#{target_shard}' is not configured in database.yml"
+        next
+      end
 
-      if website.shard_name != "default"
+      if website.shard_name != 'default'
         puts "ERROR: Website is already on shard '#{website.shard_name}'. Migration is required to move it."
         next
       end
 
-      # Update the shard name
-      website.update!(shard_name: target_shard)
+      website.update!(shard_name: target_shard.to_s)
       puts "Website #{website.id} (#{website.subdomain}) assigned to shard '#{target_shard}'."
-      puts "Note: This task only updates the pointer. If data already existed, use pwb:sharding:migrate."
+      puts 'Note: This task only updates the pointer. If data exists, use pwb:sharding:migrate.'
+    end
+
+    desc "Migrate tenant data to a target shard"
+    task :migrate, [:website_id, :target_shard] => :environment do |_, args|
+      unless args[:website_id] && args[:target_shard]
+        puts 'Usage: rake pwb:sharding:migrate[website_id,target_shard]'
+        next
+      end
+
+      website = Pwb::Website.find(args[:website_id])
+      target_shard = args[:target_shard].to_sym
+
+      migrator = Pwb::TenantShardMigrator.new(website: website, target_shard: target_shard)
+      migrator.call
+      puts "Website #{website.id} migrated to #{target_shard}."
+    rescue Pwb::TenantShardMigrator::MigrationError => e
+      puts "Migration aborted: #{e.message}"
     end
   end
 end
diff --git a/spec/services/pwb/tenant_shard_migrator_spec.rb b/spec/services/pwb/tenant_shard_migrator_spec.rb
@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+RSpec.describe Pwb::TenantShardMigrator do
+  let(:website) { create(:pwb_website, shard_name: 'default') }
+  let(:logger) { Logger.new(nil) }
+
+  describe '#call' do
+    it 'invokes insert/delete cycles and updates shard name' do
+      migrator = described_class.new(website: website, target_shard: :shard_1, logger: logger)
+
+      allow(migrator).to receive(:tenant_table_names).and_return(['pwb_contacts'])
+      allow(migrator).to receive(:fetch_rows).and_return([
+                                                          { 'id' => 42, 'website_id' => website.id, 'first_name' => 'Jane' }
+                                                        ], [])
+      allow(migrator).to receive(:ensure_no_conflicts!).and_return(true)
+
+      expect(migrator).to receive(:insert_rows).with('pwb_contacts', array_including(hash_including('id' => 42)), :shard_1)
+      expect(migrator).to receive(:delete_rows).with('pwb_contacts', [42], :default)
+
+      migrator.call
+
+      expect(website.reload.shard_name).to eq('shard_1')
+    end
+
+    it 'raises when the target shard is not configured' do
+      migrator = described_class.new(website: website, target_shard: :shard_2, logger: logger)
+      expect { migrator.call }.to raise_error(Pwb::TenantShardMigrator::MigrationError)
+    end
+  end
+end
