Merge branch 'backstage_add_type_casting_to_migration_helpers' into 'master'

Type casting for change_column_type_concurrently

See merge request gitlab-org/gitlab!34846

doc/development/what_requires_downtime.md

@@ -202,6 +202,21 @@ end
 
 And that's it, we're done!
 
+### Casting data to a new type
+
+Some type changes require casting data to a new type. For example when changing from `text` to `jsonb`.
+In this case, use the `type_cast_function` option.
+Make sure there is no bad data and the cast will always succeed. You can also provide a custom function that handles
+casting errors.
+
+Example migration:
+
+```ruby
+  def up
+    change_column_type_concurrently :users, :settings, :jsonb, type_cast_function: 'jsonb'
+  end
+```
+
 ## Changing The Schema For Large Tables
 
 While `change_column_type_concurrently` and `rename_column_concurrently` can be

lib/gitlab/database/migration_helpers.rb

@@ -477,7 +477,7 @@ module Gitlab
       #        type is used.
       # batch_column_name - option is for tables without primary key, in this
       #        case another unique integer column can be used. Example: :user_id
-      def rename_column_concurrently(table, old, new, type: nil, batch_column_name: :id)
+      def rename_column_concurrently(table, old, new, type: nil, type_cast_function: nil, batch_column_name: :id)
         unless column_exists?(table, batch_column_name)
           raise "Column #{batch_column_name} does not exist on #{table}"
         end
@@ -488,7 +488,7 @@ module Gitlab
 
         check_trigger_permissions!(table)
 
-        create_column_from(table, old, new, type: type, batch_column_name: batch_column_name)
+        create_column_from(table, old, new, type: type, batch_column_name: batch_column_name, type_cast_function: type_cast_function)
 
         install_rename_triggers(table, old, new)
       end
@@ -536,10 +536,10 @@ module Gitlab
       # table - The table containing the column.
       # column - The name of the column to change.
       # new_type - The new column type.
-      def change_column_type_concurrently(table, column, new_type)
+      def change_column_type_concurrently(table, column, new_type, type_cast_function: nil)
         temp_column = "#{column}_for_type_change"
 
-        rename_column_concurrently(table, column, temp_column, type: new_type)
+        rename_column_concurrently(table, column, temp_column, type: new_type, type_cast_function: type_cast_function)
       end
 
       # Performs cleanup of a concurrent type change.
@@ -1268,7 +1268,7 @@ into similar problems in the future (e.g. when new tables are created).
         "ON DELETE #{on_delete.upcase}"
       end
 
-      def create_column_from(table, old, new, type: nil, batch_column_name: :id)
+      def create_column_from(table, old, new, type: nil, batch_column_name: :id, type_cast_function: nil)
         old_col = column_for(table, old)
         new_type = type || old_col.type
 
@@ -1282,7 +1282,13 @@ into similar problems in the future (e.g. when new tables are created).
         # necessary since we copy over old values further down.
         change_column_default(table, new, old_col.default) unless old_col.default.nil?
 
-        update_column_in_batches(table, new, Arel::Table.new(table)[old], batch_column_name: batch_column_name)
+        old_value = Arel::Table.new(table)[old]
+
+        if type_cast_function.present?
+          old_value = Arel::Nodes::NamedFunction.new(type_cast_function, [old_value])
+        end
+
+        update_column_in_batches(table, new, old_value, batch_column_name: batch_column_name)
 
         add_not_null_constraint(table, new) unless old_col.null
 

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -690,12 +690,28 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           model.rename_column_concurrently(:users, :old, :new)
         end
 
+        context 'with existing records and type casting' do
+          let(:trigger_name) { model.rename_trigger_name(:users, :id, :new) }
+          let(:user) { create(:user) }
+
+          it 'copies the value to the new column using the type_cast_function', :aggregate_failures do
+            expect(model).to receive(:copy_indexes).with(:users, :id, :new)
+            expect(model).to receive(:add_not_null_constraint).with(:users, :new)
+            expect(model).to receive(:execute).with("UPDATE \"users\" SET \"new\" = cast_to_jsonb_with_default(\"users\".\"id\") WHERE \"users\".\"id\" >= #{user.id}")
+            expect(model).to receive(:execute).with("DROP TRIGGER IF EXISTS #{trigger_name}\nON \"users\"\n")
+            expect(model).to receive(:execute).with("CREATE TRIGGER #{trigger_name}\nBEFORE INSERT OR UPDATE\nON \"users\"\nFOR EACH ROW\nEXECUTE PROCEDURE #{trigger_name}()\n")
+            expect(model).to receive(:execute).with("CREATE OR REPLACE FUNCTION #{trigger_name}()\nRETURNS trigger AS\n$BODY$\nBEGIN\n  NEW.\"new\" := NEW.\"id\";\n  RETURN NEW;\nEND;\n$BODY$\nLANGUAGE 'plpgsql'\nVOLATILE\n")
+
+            model.rename_column_concurrently(:users, :id, :new, type_cast_function: 'cast_to_jsonb_with_default')
+          end
+        end
+
         it 'passes the batch_column_name' do
           expect(model).to receive(:column_exists?).with(:users, :other_batch_column).and_return(true)
           expect(model).to receive(:check_trigger_permissions!).and_return(true)
 
           expect(model).to receive(:create_column_from).with(
-            :users, :old, :new, type: nil, batch_column_name: :other_batch_column
+            :users, :old, :new, type: nil, batch_column_name: :other_batch_column, type_cast_function: nil
           ).and_return(true)
 
           expect(model).to receive(:install_rename_triggers).and_return(true)
@@ -703,6 +719,14 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           model.rename_column_concurrently(:users, :old, :new, batch_column_name: :other_batch_column)
         end
 
+        it 'passes the type_cast_function' do
+          expect(model).to receive(:create_column_from).with(
+            :users, :old, :new, type: nil, batch_column_name: :id, type_cast_function: 'JSON'
+          ).and_return(true)
+
+          model.rename_column_concurrently(:users, :old, :new, type_cast_function: 'JSON')
+        end
+
         it 'raises an error with invalid batch_column_name' do
           expect do
             model.rename_column_concurrently(:users, :old, :new, batch_column_name: :invalid)
@@ -866,10 +890,19 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
   describe '#change_column_type_concurrently' do
     it 'changes the column type' do
       expect(model).to receive(:rename_column_concurrently)
-        .with('users', 'username', 'username_for_type_change', type: :text)
+        .with('users', 'username', 'username_for_type_change', type: :text, type_cast_function: nil)
 
       model.change_column_type_concurrently('users', 'username', :text)
     end
+
+    context 'with type cast' do
+      it 'changes the column type with casting the value to the new type' do
+        expect(model).to receive(:rename_column_concurrently)
+          .with('users', 'username', 'username_for_type_change', type: :text, type_cast_function: 'JSON')
+
+        model.change_column_type_concurrently('users', 'username', :text, type_cast_function: 'JSON')
+      end
+    end
   end
 
   describe '#cleanup_concurrent_column_type_change' do