Add migration helpers for copying check constraints

- Add check_constraints_for(table, column, schema: nil)
  private method to Gitlab::Database::MigrationHelpers.
  Returns all the check constraints defined for a column.
- Add the copy_check_constraints(table, old, new, schema: nil)
  migration helper.
  It copies all co straints defined in column old to column new.
- Update the create_column_from migration helper to
  also copy all existing check constraints to the new column.
- Update the specs for all helpers that use create_column_from
  and add new specs for copy_check_constraints

changelogs/unreleased/247489-update-create-column-from-to-also-copy-constraints.yml

+---
+title: Add migration helpers for copying check constraints
+merge_request: 44025
+author:
+type: other

lib/gitlab/database/migration_helpers.rb

@@ -1151,6 +1151,48 @@ into similar problems in the future (e.g. when new tables are created).
         end
       end
 
+      # Copies all check constraints for the old column to the new column.
+      #
+      # table - The table containing the columns.
+      # old - The old column.
+      # new - The new column.
+      # schema - The schema the table is defined for
+      #          If it is not provided, then the current_schema is used
+      def copy_check_constraints(table, old, new, schema: nil)
+        if transaction_open?
+          raise 'copy_check_constraints can not be run inside a transaction'
+        end
+
+        unless column_exists?(table, old)
+          raise "Column #{old} does not exist on #{table}"
+        end
+
+        unless column_exists?(table, new)
+          raise "Column #{new} does not exist on #{table}"
+        end
+
+        table_with_schema = schema.present? ? "#{schema}.#{table}" : table
+
+        check_constraints_for(table, old, schema: schema).each do |check_c|
+          constraint_name = begin
+            if check_c["constraint_def"] == "(#{old} IS NOT NULL)"
+              not_null_constraint_name(table_with_schema, new)
+            elsif check_c["constraint_def"].start_with? "(char_length(#{old}) <="
+              text_limit_name(table_with_schema, new)
+            else
+              check_constraint_name(table_with_schema, new, 'copy_check_constraint')
+            end
+          end
+
+          add_check_constraint(
+            table_with_schema,
+            check_c["constraint_def"].gsub(old.to_s, new.to_s),
+            constraint_name,
+            validate: true
+          )
+        end
+      end
+
       # Migration Helpers for adding limit to text columns
       def add_text_limit(table, column, limit, constraint_name: nil, validate: true)
         add_check_constraint(
@@ -1278,6 +1320,37 @@ into similar problems in the future (e.g. when new tables are created).
         end
       end
 
+      # Returns an ActiveRecord::Result containing the check constraints
+      # defined for the given column.
+      #
+      # If the schema is not provided, then the current_schema is used
+      def check_constraints_for(table, column, schema: nil)
+        check_sql = <<~SQL
+          SELECT
+            ccu.table_schema as schema_name,
+            ccu.table_name as table_name,
+            ccu.column_name as column_name,
+            con.conname as constraint_name,
+            con.consrc as constraint_def
+          FROM pg_catalog.pg_constraint con
+            INNER JOIN pg_catalog.pg_class rel
+              ON rel.oid = con.conrelid
+            INNER JOIN pg_catalog.pg_namespace nsp
+              ON nsp.oid = con.connamespace
+            INNER JOIN information_schema.constraint_column_usage ccu
+              ON con.conname = ccu.constraint_name
+                     AND nsp.nspname = ccu.constraint_schema
+                     AND rel.relname = ccu.table_name
+          WHERE  nsp.nspname = #{connection.quote(schema.presence || current_schema)}
+            AND rel.relname = #{connection.quote(table)}
+            AND ccu.column_name = #{connection.quote(column)}
+            AND con.contype = 'c'
+          ORDER BY constraint_name
+        SQL
+
+        connection.exec_query(check_sql)
+      end
+
       def statement_timeout_disabled?
         # This is a string of the form "100ms" or "0" when disabled
         connection.select_value('SHOW statement_timeout') == "0"
@@ -1357,6 +1430,7 @@ into similar problems in the future (e.g. when new tables are created).
 
         copy_indexes(table, old, new)
         copy_foreign_keys(table, old, new)
+        copy_check_constraints(table, old, new)
       end
 
       def validate_timestamp_column_name!(column_name)

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -699,6 +699,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
           expect(model).to receive(:copy_indexes).with(:users, :old, :new)
           expect(model).to receive(:copy_foreign_keys).with(:users, :old, :new)
+          expect(model).to receive(:copy_check_constraints).with(:users, :old, :new)
 
           model.rename_column_concurrently(:users, :old, :new)
         end
@@ -761,6 +762,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
             expect(model).to receive(:change_column_default)
               .with(:users, :new, old_column.default)
 
+            expect(model).to receive(:copy_check_constraints)
+              .with(:users, :old, :new)
+
             model.rename_column_concurrently(:users, :old, :new)
           end
         end
@@ -856,6 +860,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
         expect(model).to receive(:copy_indexes).with(:users, :new, :old)
         expect(model).to receive(:copy_foreign_keys).with(:users, :new, :old)
+        expect(model).to receive(:copy_check_constraints).with(:users, :new, :old)
 
         model.undo_cleanup_concurrent_column_rename(:users, :old, :new)
       end
@@ -894,6 +899,9 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:change_column_default)
             .with(:users, :old, new_column.default)
 
+          expect(model).to receive(:copy_check_constraints)
+            .with(:users, :new, :old)
+
           model.undo_cleanup_concurrent_column_rename(:users, :old, :new)
         end
       end
@@ -2172,6 +2180,103 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#copy_check_constraints' do
+    context 'inside a transaction' do
+      it 'raises an error' do
+        expect(model).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          model.copy_check_constraints(:test_table, :old_column, :new_column)
+        end.to raise_error(RuntimeError)
+      end
+    end
+
+    context 'outside a transaction' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+        allow(model).to receive(:column_exists?).and_return(true)
+      end
+
+      let(:old_column_constraints) do
+        [
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'check_d7d49d475d',
+            'constraint_def' => '(old_column IS NOT NULL)'
+          },
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'check_48560e521e',
+            'constraint_def' => '(char_length(old_column) <= 255)'
+          },
+          {
+            'schema_name' => 'public',
+            'table_name' => 'test_table',
+            'column_name' => 'old_column',
+            'constraint_name' => 'custom_check_constraint',
+            'constraint_def' => '((old_column IS NOT NULL) AND (another_column IS NULL))'
+          }
+        ]
+      end
+
+      it 'copies check constraints from one column to another' do
+        allow(model).to receive(:check_constraints_for)
+        .with(:test_table, :old_column, schema: nil)
+          .and_return(old_column_constraints)
+
+        allow(model).to receive(:not_null_constraint_name).with(:test_table, :new_column)
+          .and_return('check_1')
+
+        allow(model).to receive(:text_limit_name).with(:test_table, :new_column)
+          .and_return('check_2')
+
+        allow(model).to receive(:check_constraint_name)
+          .with(:test_table, :new_column, 'copy_check_constraint')
+          .and_return('check_3')
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '(new_column IS NOT NULL)',
+            'check_1',
+            validate: true
+          ).once
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '(char_length(new_column) <= 255)',
+            'check_2',
+            validate: true
+          ).once
+
+        expect(model).to receive(:add_check_constraint)
+          .with(
+            :test_table,
+            '((new_column IS NOT NULL) AND (another_column IS NULL))',
+            'check_3',
+            validate: true
+          ).once
+
+        model.copy_check_constraints(:test_table, :old_column, :new_column)
+      end
+
+      it 'does nothing if there are no constraints defined for the old column' do
+        allow(model).to receive(:check_constraints_for)
+        .with(:test_table, :old_column, schema: nil)
+          .and_return([])
+
+        expect(model).not_to receive(:add_check_constraint)
+
+        model.copy_check_constraints(:test_table, :old_column, :new_column)
+      end
+    end
+  end
+
   describe '#add_text_limit' do
     context 'when it is called with the default options' do
       it 'calls add_check_constraint with an infered constraint name and validate: true' do