ruby_llm 2.0.0.rc2 → 2.0.0.rc4

data/lib/generators/ruby_llm/upgrade/online_copy_migration/verification.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+require_relative '../legacy_content_sql'
+
+module RubyLLM
+  module Generators
+    class OnlineCopyMigration < UpgradeMigration
+      class Verification # :nodoc: all
+        def initialize(connection:, settings:)
+          @connection = connection
+          @settings = settings
+        end
+
+        def verify
+          verify_chats
+          verify_content
+          verify_usages
+          verify_tools
+        end
+
+        private
+
+        def verify_chats
+          mismatch = %w[provider model_id].map do |column|
+            different("copied.#{q(column)}", "original.#{q(column)}", string: true)
+          end
+          reject_mismatch('chat model', <<~SQL)
+            SELECT c.#{pk('chat')} FROM #{table('chat')} c
+            LEFT JOIN #{table('model')} original ON original.#{pk('model')} = c.#{q(@settings.fetch('model_foreign_key'))}
+            LEFT JOIN #{qt(:ruby_llm_models)} copied ON copied.#{pk('model')} = c.ruby_llm_model_id
+            WHERE c.ruby_llm_version = 1
+              AND (original.#{pk('model')} IS NULL OR copied.#{pk('model')} IS NULL OR #{mismatch.join(' OR ')})
+          SQL
+        end
+
+        def verify_content
+          raw = source(:content_raw)
+          rendered = LegacyContentSQL.new(@connection).render(content: "m.#{q(:content)}", raw:)
+          fields = { ruby_llm_content: rendered, raw_content: raw }
+          failed = fields.map do |column, expected|
+            different("m.#{q(column)}", expected, json: column == :raw_content, string: true)
+          end
+          reject_mismatch('message content', "SELECT m.#{pk('message')} #{message_from} " \
+                                             "WHERE c.ruby_llm_version = 1 AND (#{failed.join(' OR ')})")
+        end
+
+        def verify_usages
+          fields = usage_fields
+          eligible = fields.slice(*%i[input_tokens output_tokens cache_read_tokens cache_write_tokens thinking_tokens
+                                      input_cost output_cost cache_read_cost cache_write_cost thinking_cost total_cost])
+                           .values.map { |value| "#{value} IS NOT NULL" }
+          eligible << "m.#{q(:role)} = 'assistant'"
+          strings = %i[chat_type message_type provider model operation status]
+          failed = fields.map do |column, expected|
+            different("u.#{q(column)}", expected, string: strings.include?(column))
+          end
+          reject_mismatch('usage entry', <<~SQL)
+            SELECT m.#{pk('message')} #{message_from}
+            LEFT JOIN #{table('model')} cm ON cm.#{pk('model')} = c.#{q(@settings.fetch('model_foreign_key'))}
+            #{message_model_join}
+            LEFT JOIN #{qt(:ruby_llm_usages)} u
+              ON u.legacy_key = #{copied_key(:ruby_llm_usages, "m.#{pk('message')}")}
+            WHERE c.ruby_llm_version = 1 AND (#{eligible.join(' OR ')})
+              AND (u.id IS NULL OR #{failed.join(' OR ')})
+          SQL
+        end
+
+        def usage_fields
+          fields = {
+            input_tokens: source(:input_tokens), output_tokens: source(:output_tokens),
+            cache_read_tokens: coalesce(source(:cache_read_tokens), source(:cached_tokens)),
+            cache_write_tokens: coalesce(source(:cache_write_tokens), source(:cache_creation_tokens)),
+            thinking_tokens: source(:thinking_tokens)
+          }
+          %w[input output cache_read cache_write thinking].each { |key| fields[:"#{key}_cost"] = cost(key) }
+          fields[:total_cost] = coalesce(source(:total_cost), cost('total'))
+          provider, model = model_identity
+          fields.merge(chat_type: value(@settings.fetch('chat_class')), chat_id: "c.#{pk('chat')}",
+                       message_type: value(@settings.fetch('message_class')), message_id: "m.#{pk('message')}",
+                       provider:, model:, operation: "'chat'", status: "'succeeded'")
+        end
+
+        def model_identity
+          key = @settings.fetch('model_foreign_key')
+          if message_columns[key]&.type == :string
+            ["COALESCE(#{source(:provider)}, cm.provider)", "COALESCE(#{source(key)}, cm.model_id)"]
+          elsif message_columns.key?(key)
+            ['COALESCE(mm.provider, cm.provider)', 'COALESCE(mm.model_id, cm.model_id)']
+          else
+            ['cm.provider', 'cm.model_id']
+          end
+        end
+
+        def message_model_join
+          key = @settings.fetch('model_foreign_key')
+          return '' if !message_columns.key?(key) || message_columns[key].type == :string
+
+          "LEFT JOIN #{table('model')} mm ON mm.#{pk('model')} = m.#{q(key)}"
+        end
+
+        def cost(key)
+          return 'NULL' unless message_columns.key?('cost_details')
+
+          details = source(:cost_details)
+          case @connection.adapter_name
+          when 'PostgreSQL' then "NULLIF(#{details}::jsonb ->> '#{key}', '')::numeric"
+          when 'Mysql2'
+            "CAST(NULLIF(JSON_UNQUOTE(JSON_EXTRACT(#{details}, '$.#{key}')), 'null') AS DECIMAL(16, 10))"
+          else "json_extract(#{details}, '$.#{key}')"
+          end
+        end
+
+        def verify_tools
+          result_type = "CASE WHEN r.#{pk('message')} IS NOT NULL THEN #{value(@settings.fetch('message_class'))} END"
+          fields = { message_id: "s.#{q(@settings.fetch('message_foreign_key'))}",
+                     message_type: value(@settings.fetch('message_class')), result_id: "r.#{pk('message')}",
+                     result_type:,
+                     tool_call_id: 's.tool_call_id', name: 's.name', arguments: 's.arguments' }
+          strings = %i[message_type result_type tool_call_id name arguments]
+          failed = fields.map do |column, expected|
+            different("t.#{q(column)}", expected, json: column == :arguments, string: strings.include?(column))
+          end
+          reject_mismatch('tool call', <<~SQL)
+            SELECT s.#{pk('tool_call')} FROM #{table('tool_call')} s
+            LEFT JOIN #{table('message')} m ON m.#{pk('message')} = s.#{q(@settings.fetch('message_foreign_key'))}
+            LEFT JOIN #{table('chat')} c ON c.#{pk('chat')} = m.#{q(@settings.fetch('chat_foreign_key'))}
+            LEFT JOIN #{table('message')} r ON r.#{q(@settings.fetch('tool_call_foreign_key'))} = s.#{pk('tool_call')}
+            LEFT JOIN #{qt(:ruby_llm_tool_calls)} t
+              ON t.legacy_key = #{copied_key(:ruby_llm_tool_calls, "s.#{pk('tool_call')}")}
+            WHERE c.ruby_llm_version = 1 AND (t.id IS NULL OR #{failed.join(' OR ')})
+          SQL
+        end
+
+        def message_from
+          "FROM #{table('message')} m JOIN #{table('chat')} c " \
+            "ON c.#{pk('chat')} = m.#{q(@settings.fetch('chat_foreign_key'))}"
+        end
+
+        def different(left, right, json: false, string: false)
+          left, right = comparable([left, right], json:, string:)
+          case @connection.adapter_name
+          when 'PostgreSQL' then "#{left} IS DISTINCT FROM #{right}"
+          when 'Mysql2' then "NOT (#{left} <=> #{right})"
+          else "#{left} IS NOT #{right}"
+          end
+        end
+
+        def comparable(expressions, json:, string:)
+          case @connection.adapter_name
+          when 'PostgreSQL'
+            json ? expressions.map { |expression| "(#{expression})::jsonb" } : expressions
+          when 'Mysql2'
+            string ? expressions.map { |expression| "CAST(#{expression} AS BINARY)" } : expressions
+          else
+            json ? expressions.map { |expression| "json(#{expression})" } : expressions
+          end
+        end
+
+        def reject_mismatch(label, sql)
+          id = @connection.select_value("#{sql.strip} LIMIT 1")
+          raise "Record #{id} did not preserve its #{label}" if id
+        end
+
+        def text(expression) = mysql? ? "CAST(#{expression} AS BINARY)" : "CAST(#{expression} AS TEXT)"
+
+        def copied_key(table, expression)
+          return text(expression) unless mysql?
+
+          collation = @connection.columns(table).find { |column| column.name == 'legacy_key' }.collation
+          charset = collation.split('_').first
+          "CAST(#{expression} AS CHAR CHARACTER SET #{q(charset)}) COLLATE #{q(collation)}"
+        end
+
+        def coalesce(*expressions)
+          values = expressions.reject { |expression| expression == 'NULL' }
+          return 'NULL' if values.empty?
+          return values.first if values.one?
+
+          "COALESCE(#{values.join(', ')})"
+        end
+
+        def source(column) = message_columns.key?(column.to_s) ? "m.#{q(column)}" : 'NULL'
+
+        def message_columns
+          @message_columns ||= @connection.columns(@settings.fetch('message_table')).index_by(&:name)
+        end
+
+        def mysql? = @connection.adapter_name == 'Mysql2'
+        def table(kind) = qt(@settings.fetch("#{kind}_table"))
+        def pk(kind) = q(@connection.primary_key(@settings.fetch("#{kind}_table")))
+        def value(value) = @connection.quote(value)
+        def q(value) = @connection.quote_column_name(value)
+        def qt(value) = @connection.quote_table_name(value)
+      end
+    end
+  end
+end

data/lib/generators/ruby_llm/upgrade/online_copy_migration.rb

@@ -0,0 +1,174 @@
+# frozen_string_literal: true
+
+require_relative 'upgrade_migration'
+require_relative 'online_copy_migration/journal'
+require_relative 'online_copy_migration/data'
+
+module RubyLLM
+  module Generators
+    class OnlineCopyMigration < UpgradeMigration # :nodoc: all
+      def online? = true
+
+      def prepare(settings)
+        with_migration_lock do
+          super(settings.merge(online: true), &nil)
+          unless @connection.column_exists?(TABLE, :epoch)
+            @connection.add_column(TABLE, :epoch, :bigint, null: false, default: 0)
+          end
+          yield if block_given?
+          data.prepare
+          journal.install
+        end
+      end
+
+      def verify_settings(settings)
+        super(settings.merge(online: true))
+      end
+
+      def copy_table(source, target)
+        create_copy_table(source, target) unless @connection.table_exists?(target)
+      end
+
+      def copy_model_references
+        nil
+      end
+
+      def backfill
+        with_migration_lock do
+          state = states.first!
+          unless state.active_version == 1 && %w[preparing active].include?(state.status)
+            raise 'Online copy backfill requires RubyLLM 1.16 to be active'
+          end
+          raise 'Use resume after rolling back a finished upgrade' if data.finished?
+
+          data.backfill
+          data.catch_up(passes: 2)
+          state.update!(needs_reconcile: false)
+        end
+      end
+
+      def finish(discard_incomplete_tool_calls: false)
+        with_migration_lock do
+          state = states.first!
+          return if state.active_version == 2 && state.status == 'active' && data.finished?
+
+          raise 'Run the copy backfill before finish' unless data.completed?
+
+          pause(state, from: 1, statuses: %w[preparing active finishing])
+          if discard_incomplete_tool_calls
+            discarded = data.discard_incomplete_tool_calls
+            ::ActiveRecord::Migration.say "Discarded #{discarded} incomplete legacy tool calls"
+          end
+          data.catch_up
+          data.verify
+          verify_completed_work
+          yield
+          journal.install
+          state.update!(active_version: 2, status: 'active', needs_reconcile: false)
+        end
+      end
+
+      def rollback
+        with_migration_lock do
+          state = states.first!
+          unless data.finished?
+            raise 'The unfinished upgrade is not running on 1.16' unless state.active_version == 1
+
+            state.update!(status: 'preparing', needs_reconcile: true)
+            return
+          end
+          state.with_lock do
+            require_active_version(state, 2)
+            verify_completed_work
+            state.update!(active_version: 1, epoch: state.epoch + 1)
+          end
+        end
+      end
+
+      def resume
+        with_migration_lock do
+          raise 'Finish the initial copy migrations before resuming 2.0' unless data.finished?
+
+          state = states.first!
+          pause(state, from: 1, statuses: %w[active finishing])
+          data.catch_up
+          data.verify
+          verify_completed_work
+          state.update!(active_version: 2, status: 'active', needs_reconcile: false)
+        end
+      end
+
+      def finalize
+        with_migration_lock { super }
+      end
+
+      def cleanup
+        with_migration_lock do
+          verify_cleanup
+          journal.remove
+          yield if block_given?
+          data.cleanup
+          super
+        end
+      end
+
+      private
+
+      def pause(state, from:, statuses:)
+        state.with_lock do
+          unless state.active_version == from && statuses.include?(state.status)
+            raise 'The copy upgrade is not ready for this version switch'
+          end
+
+          state.update!(status: 'finishing', epoch: state.epoch + 1) unless state.status == 'finishing'
+        end
+      end
+
+      def data
+        @data ||= Data.new(connection: @connection, settings: configuration)
+      end
+
+      def journal
+        @journal ||= Journal.new(connection: @connection, settings: configuration)
+      end
+
+      def with_migration_lock(&)
+        return with_mysql_lock(&) if @connection.adapter_name == 'Mysql2'
+        return with_sqlite_lock(&) if @connection.adapter_name == 'SQLite'
+
+        key = "hashtextextended(current_database() || '.' || current_schema() || '.ruby_llm_upgrade', 0)"
+        locked = @connection.select_value("SELECT pg_try_advisory_lock(#{key})")
+        raise 'Another RubyLLM copy migration is running' unless locked
+
+        begin
+          yield
+        ensure
+          @connection.execute("SELECT pg_advisory_unlock(#{key})")
+        end
+      end
+
+      def with_mysql_lock
+        key = @connection.quote("ruby_llm_upgrade:#{@connection.current_database}".slice(0, 64))
+        locked = @connection.select_value("SELECT GET_LOCK(#{key}, 0)")
+        raise 'Another RubyLLM copy migration is running' unless locked == 1
+
+        begin
+          yield
+        ensure
+          @connection.execute("SELECT RELEASE_LOCK(#{key})")
+        end
+      end
+
+      def with_sqlite_lock
+        database = @connection.pool.db_config.database
+        return yield if database == ':memory:'
+
+        File.open("#{database}.ruby_llm_upgrade.lock", File::RDWR | File::CREAT, 0o600) do |file|
+          raise 'Another RubyLLM copy migration is running' unless file.flock(File::LOCK_EX | File::LOCK_NB)
+
+          yield
+        end
+      end
+    end
+  end
+end

data/lib/generators/ruby_llm/upgrade/templates/backfill_v2_data.rb.tt

@@ -1,33 +1,34 @@
 # frozen_string_literal: true
-<% if copy_mode? -%>
 
+<% if copy_mode? -%>
 require 'generators/ruby_llm/upgrade/upgrade_migration'
+<% else -%>
+require 'generators/ruby_llm/upgrade/legacy_content_sql'
 <% end -%>
 
 class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
   disable_ddl_transaction!
 
+<% unless copy_mode? -%>
   BATCH_SIZE = 10_000
+<% end -%>
   PROGRESS_TABLE = :ruby_llm_v2_backfills
 
   def up
 <% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.verify_settings(<%= copy_upgrade_settings.inspect %>)
+    RubyLLM::Generators::UpgradeMigration.for.verify_settings(<%= copy_upgrade_settings.inspect %>)
+    verify_upgrade_in_progress
+    RubyLLM::Generators::UpgradeMigration.for.backfill
 <% else -%>
     raise 'Generate this migration with --mode copy' if table_exists?(:ruby_llm_v2_upgrades)
-<% end -%>
     verify_upgrade_in_progress
     create_progress_table
     backfill_required_defaults
-<% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.backfill
-    %w[message_content tool_results usages].each { |task| mark_completed(task) }
-<% else -%>
     backfill_message_content
     backfill_tool_results
     backfill_usages
-<% end -%>
     verify_backfills
+<% end -%>
   end
 
   def down
@@ -47,6 +48,7 @@ class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migratio
     raise 'The RubyLLM 2.0 upgrade is already finished. Do not run the backfill again.'
   end
 
+<% unless copy_mode? -%>
   def backfill_required_defaults
     {
       <%= chat_table_name %>: {cancelled: false},
@@ -69,18 +71,16 @@ class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migratio
       content = message_value(:content)
 <% if postgresql? -%>
       structured_content = "#{content_raw}::jsonb"
-      rendered_content = "#{content_raw}::text"
 <% elsif mysql? -%>
       structured_content = content_raw
-      rendered_content = "CAST(#{content_raw} AS CHAR)"
 <% else -%>
       structured_content = content_raw
-      rendered_content = content_raw
 <% end -%>
+      rendered_content = RubyLLM::Generators::LegacyContentSQL.new(connection).render(content: content, raw: content_raw)
       execute <<~SQL
         UPDATE #{quote_table(:<%= message_table_name %>)} AS legacy_messages
            SET #{quote_column(:raw_content)} = COALESCE(#{raw_content}, #{structured_content}),
-               #{quote_column(:content)} = COALESCE(#{content}, #{rendered_content})
+               #{quote_column(:content)} = #{rendered_content}
          WHERE #{range}
            AND #{content_raw} IS NOT NULL
       SQL
@@ -187,6 +187,10 @@ class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migratio
   end
 
   def verify_backfills
+<% if postgresql? -%>
+    tables = %i[<%= chat_table_name %> <%= message_table_name %> ruby_llm_models ruby_llm_tool_calls ruby_llm_usages]
+    with_upgrade_safety { execute "ANALYZE #{tables.map { |table| quote_table(table) }.join(', ')}" }
+<% end -%>
     verify_message_content
     verify_tool_results
     verify_usages
@@ -300,6 +304,7 @@ class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migratio
 
   def usage_candidate_sql(conditions, joins, range)
     usages = quote_table(:ruby_llm_usages)
+    usage_range = range.gsub(message_value(connection.primary_key(:<%= message_table_name %>)), 'existing_usages.message_id')
     <<~SQL
       FROM #{quote_table(:<%= message_table_name %>)} legacy_messages
       #{joins}
@@ -309,6 +314,7 @@ class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migratio
           SELECT 1
             FROM #{usages} existing_usages
            WHERE existing_usages.message_type = #{connection.quote('<%= message_model_name %>')}
+             AND (#{usage_range})
              AND existing_usages.message_id = #{message_value(connection.primary_key(:<%= message_table_name %>))}
              AND existing_usages.operation = 'chat'
              AND existing_usages.status = 'succeeded'
@@ -450,6 +456,7 @@ class <%= backfill_migration_class_name %> < ActiveRecord::Migration<%= migratio
     yield
   end
 
+<% end -%>
   def progress_records = migration_record(PROGRESS_TABLE)
 
   def migration_record(table)

data/lib/generators/ruby_llm/upgrade/templates/cleanup_v2_upgrade.rb.tt

@@ -11,15 +11,14 @@ class <%= cleanup_migration_class_name %> < ActiveRecord::Migration<%= migration
 
   def up
 <% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.verify_settings(<%= copy_upgrade_settings.inspect %>)
-    RubyLLM::Generators::UpgradeMigration.new.verify_cleanup
+    RubyLLM::Generators::UpgradeMigration.for.verify_settings(<%= copy_upgrade_settings.inspect %>)
+    RubyLLM::Generators::UpgradeMigration.for.verify_cleanup
 <% else -%>
     raise 'Generate cleanup with --mode copy for this database' if table_exists?(:ruby_llm_v2_upgrades)
 <% end -%>
     unless table_exists?(PROGRESS_TABLE)
 <% if copy_mode? -%>
-      remove_legacy_message_columns
-      RubyLLM::Generators::UpgradeMigration.new.cleanup
+      RubyLLM::Generators::UpgradeMigration.for.cleanup { remove_legacy_message_columns }
       return
 <% else -%>
       return unless legacy_columns.any?
@@ -31,9 +30,10 @@ class <%= cleanup_migration_class_name %> < ActiveRecord::Migration<%= migration
       raise 'Run <%= finish_migration_class_name %> before removing the legacy message columns'
     end
 
-    remove_legacy_message_columns
 <% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.cleanup
+    RubyLLM::Generators::UpgradeMigration.for.cleanup { remove_legacy_message_columns }
+<% else -%>
+    remove_legacy_message_columns
 <% end -%>
     drop_table PROGRESS_TABLE
   end

data/lib/generators/ruby_llm/upgrade/templates/finish_v2_upgrade.rb.tt

@@ -12,20 +12,26 @@ class <%= finish_migration_class_name %> < ActiveRecord::Migration<%= migration_
 
   def up
 <% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.verify_settings(<%= copy_upgrade_settings.inspect %>)
-    RubyLLM::Generators::UpgradeMigration.new.backfill(force: false)
+    RubyLLM::Generators::UpgradeMigration.for.verify_settings(<%= copy_upgrade_settings.inspect %>)
+    RubyLLM::Generators::UpgradeMigration.for.finish<% if discard_incomplete_tool_calls? %>(discard_incomplete_tool_calls: true)<% end %> do
+      verify_completed_backfills
+      enforce_required_defaults
+      relax_legacy_message_constraints
+      mark_finished
+    end
 <% else -%>
     raise 'Generate this migration with --mode copy' if table_exists?(:ruby_llm_v2_upgrades)
-<% end -%>
     verify_completed_backfills
+<% if postgresql? -%>
+    tables = %i[<%= chat_table_name %> <%= message_table_name %> ruby_llm_models ruby_llm_tool_calls ruby_llm_usages]
+    with_upgrade_safety { execute "ANALYZE #{tables.map { |table| quote_table(table) }.join(', ')}" }
+<% end -%>
     verify_message_content
     verify_tool_results
     verify_usages
     enforce_required_defaults
     relax_legacy_message_constraints
     mark_finished
-<% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.activate
 <% end -%>
   end
 
@@ -36,6 +42,17 @@ class <%= finish_migration_class_name %> < ActiveRecord::Migration<%= migration_
   private
 
   def relax_legacy_message_constraints
+<% if copy_mode? -%>
+    model_column = :<%= v1_model_foreign_key %>
+    if column_definition(:<%= chat_table_name %>, model_column)&.null == false
+      with_upgrade_safety { change_column_null :<%= chat_table_name %>, model_column, true }
+    end
+    connection.foreign_keys(:<%= v1_tool_call_table_name %>).each do |key|
+      next unless key.column == '<%= message_foreign_key %>'
+
+      remove_foreign_key :<%= v1_tool_call_table_name %>, column: key.column
+    end
+<% end -%>
     %i[<%= v1_model_foreign_key %> <%= v1_tool_call_foreign_key %>].each do |column|
       next unless column_definition(:<%= message_table_name %>, column)&.null == false
 
@@ -51,6 +68,15 @@ class <%= finish_migration_class_name %> < ActiveRecord::Migration<%= migration_
   end
 
   def enforce_required_defaults
+    {
+      <%= chat_table_name %>: {cancelled: false},
+      <%= message_table_name %>: {cache_until_here: false},
+      ruby_llm_tool_calls: {message_type: '<%= message_model_name %>'}
+    }.each do |table, attributes|
+      migration_record(table).where(attributes.transform_values { nil }).in_batches(of: 10_000) do |batch|
+        batch.update_all(attributes)
+      end
+    end
 <% if copy_mode? -%>
     enforce_not_null(:<%= chat_table_name %>, :ruby_llm_model_id)
 <% end -%>
@@ -87,6 +113,7 @@ class <%= finish_migration_class_name %> < ActiveRecord::Migration<%= migration_
     raise "RubyLLM 2.0 backfills are incomplete: #{missing.join(', ')}"
   end
 
+<% unless copy_mode? -%>
   def verify_message_content
     return unless column_exists?(:<%= message_table_name %>, :content_raw)
 
@@ -194,17 +221,22 @@ class <%= finish_migration_class_name %> < ActiveRecord::Migration<%= migration_
 <% end -%>
   end
 
+<% end -%>
   def with_upgrade_safety(&)
     return safety_assured(&) if respond_to?(:safety_assured, true)
 
     yield
   end
 
-  def message_value(column) = "legacy_messages.#{quote_column(column)}"
   def column_definition(table, column) = connection.columns(table).find { |candidate| candidate.name == column.to_s }
+<% unless copy_mode? -%>
+  def message_value(column) = "legacy_messages.#{quote_column(column)}"
   def quoted_primary_key(table) = quote_column(connection.primary_key(table))
   def quote_table(table) = connection.quote_table_name(table)
+<% end -%>
+<% if postgresql? || !copy_mode? -%>
   def quote_column(column) = connection.quote_column_name(column)
+<% end -%>
 
   def migration_record(table)
     Class.new(ActiveRecord::Base) do

data/lib/generators/ruby_llm/upgrade/templates/prepare_v2_upgrade.rb.tt

@@ -9,16 +9,11 @@ class <%= prepare_migration_class_name %> < ActiveRecord::Migration<%= migration
 
   def up
 <% if copy_mode? -%>
-    RubyLLM::Generators::UpgradeMigration.new.prepare(<%= copy_upgrade_settings.inspect %>)
+    RubyLLM::Generators::UpgradeMigration.for.prepare(<%= copy_upgrade_settings.inspect %>) { prepare_schema }
 <% else -%>
     raise 'This database uses --mode copy' if table_exists?(:ruby_llm_v2_upgrades)
+    prepare_schema
 <% end -%>
-    validate_upgrade
-    move_models
-    add_chat_and_message_columns
-    move_tool_calls
-    create_usages
-    create_batches
   end
 
   def down
@@ -27,6 +22,15 @@ class <%= prepare_migration_class_name %> < ActiveRecord::Migration<%= migration
 
   private
 
+  def prepare_schema
+    validate_upgrade
+    move_models
+    add_chat_and_message_columns
+    move_tool_calls
+    create_usages
+    create_batches
+  end
+
   def validate_upgrade
     %i[<%= chat_table_name %> <%= message_table_name %>].each do |table|
       raise "Expected #{table} to exist before upgrading RubyLLM" unless table_exists?(table)
@@ -247,8 +251,7 @@ class <%= prepare_migration_class_name %> < ActiveRecord::Migration<%= migration
     unless column_exists?(table, :ruby_llm_model_id)
       add_column table, :ruby_llm_model_id, reference_type_for(:ruby_llm_models)
     end
-    RubyLLM::Generators::UpgradeMigration.new.copy_model_references
-    with_upgrade_safety { change_column_null table, column, true } unless column_definition(table, column).null
+    RubyLLM::Generators::UpgradeMigration.for.copy_model_references
 <% else -%>
     with_upgrade_safety { rename_column table, column, :ruby_llm_model_id } if column != :ruby_llm_model_id
 <% end -%>
@@ -288,7 +291,9 @@ class <%= prepare_migration_class_name %> < ActiveRecord::Migration<%= migration
   end
 
   def ensure_boolean_column(table, column)
-    add_column table, column, :boolean unless column_exists?(table, column)
+    unless column_exists?(table, column)
+      with_upgrade_safety { add_column table, column, :boolean, default: false }
+    end
     return if ActiveRecord::Type::Boolean.new.cast(column_definition(table, column).default) == false
 
     with_upgrade_safety { change_column_default table, column, false }
@@ -296,9 +301,6 @@ class <%= prepare_migration_class_name %> < ActiveRecord::Migration<%= migration
 
   def move_tool_calls
     move_table(:<%= v1_tool_call_table_name %>, :ruby_llm_tool_calls)
-<% if copy_mode? -%>
-    remove_column_foreign_keys(:<%= v1_tool_call_table_name %>, :<%= message_foreign_key %>)
-<% end -%>
     table = :ruby_llm_tool_calls
 <% if postgresql? -%>
     normalize_postgresql_json_columns(table, %i[arguments])
@@ -524,7 +526,7 @@ class <%= prepare_migration_class_name %> < ActiveRecord::Migration<%= migration
 
   def move_table(source, target)
 <% if copy_mode? -%>
-    with_upgrade_safety { RubyLLM::Generators::UpgradeMigration.new.copy_table(source, target) }
+    with_upgrade_safety { RubyLLM::Generators::UpgradeMigration.for.copy_table(source, target) }
 <% else -%>
     return if source == target || !table_exists?(source)