Add migration helper to create table with CHECKs

Add a new migration helper that allows the caller to define CHECK
constraints in the table definition block. This can be used to avoid
having to `disable_ddl_transaction!` and add the check constraints using
the existing helper, which cannot be used inside a transaction block.
Also provides a convenience method to more easily add CHECK constraints
for text limits.

doc/development/database/strings_and_the_text_data_type.md

@@ -11,11 +11,11 @@ info: To determine the technical writer assigned to the Stage/Group associated w
 When adding new columns that will be used to store strings or other textual information:
 
 1. We always use the `text` data type instead of the `string` data type.
-1. `text` columns should always have a limit set by using the `add_text_limit` migration helper.
+1. `text` columns should always have a limit set, either by using the `create_table_with_constraints` helper
+when creating a table, or by using the `add_text_limit` when altering an existing table.
 
-The `text` data type can not be defined with a limit, so `add_text_limit` is enforcing that by
-adding a [check constraint](https://www.postgresql.org/docs/11/ddl-constraints.html) on the
-column and then validating it at a followup step.
+The `text` data type can not be defined with a limit, so `create_table_with_constraints` and `add_text_limit` enforce
+that by adding a [check constraint](https://www.postgresql.org/docs/11/ddl-constraints.html) on the column.
 
 ## Background information
 
@@ -48,20 +48,15 @@ class CreateDbGuides < ActiveRecord::Migration[6.0]
 
   DOWNTIME = false
 
-  disable_ddl_transaction!
-
   def up
-    unless table_exists?(:db_guides)
-      create_table :db_guides do |t|
+    create_table_with_constraints :db_guides do |t|
       t.bigint :stars, default: 0, null: false
       t.text :title
       t.text :notes
-      end
-    end
 
-    # The following add the constraints and validate them immediately (no data in the table)
-    add_text_limit :db_guides, :title, 128
-    add_text_limit :db_guides, :notes, 1024
+      t.text_limit :title, 128
+      t.text_limit :notes, 1024
+    end
   end
 
   def down
@@ -71,12 +66,8 @@ class CreateDbGuides < ActiveRecord::Migration[6.0]
 end
 ```
 
-Adding a check constraint requires an exclusive lock while the `ALTER TABLE` that adds is running.
-As we don't want the exclusive lock to be held for the duration of a transaction, `add_text_limit`
-must always run in a migration with `disable_ddl_transaction!`.
-
-Also, note that we have to add a check that the table exists so that the migration can be repeated
-in case of a failure.
+Note that the `create_table_with_constraints` helper uses the `with_lock_retries` helper
+internally, so we don't need to manually wrap the method call in the migration.
 
 ## Add a text column to an existing table
 

lib/gitlab/database/migration_helpers.rb

@@ -70,6 +70,61 @@ module Gitlab
         end
       end
 
+      #
+      # Creates a new table, optionally allowing the caller to add check constraints to the table.
+      # Aside from that addition, this method should behave identically to Rails' `create_table` method.
+      #
+      # Example:
+      #
+      #     create_table_with_constraints :some_table do |t|
+      #       t.integer :thing, null: false
+      #       t.text :other_thing
+      #
+      #       t.check_constraint :thing_is_not_null, 'thing IS NOT NULL'
+      #       t.text_limit :other_thing, 255
+      #     end
+      #
+      # See Rails' `create_table` for more info on the available arguments.
+      def create_table_with_constraints(table_name, **options, &block)
+        helper_context = self
+        check_constraints = []
+
+        with_lock_retries do
+          create_table(table_name, **options) do |t|
+            t.define_singleton_method(:check_constraint) do |name, definition|
+              helper_context.send(:validate_check_constraint_name!, name) # rubocop:disable GitlabSecurity/PublicSend
+
+              check_constraints << { name: name, definition: definition }
+            end
+
+            t.define_singleton_method(:text_limit) do |column_name, limit, name: nil|
+              # rubocop:disable GitlabSecurity/PublicSend
+              name = helper_context.send(:text_limit_name, table_name, column_name, name: name)
+              helper_context.send(:validate_check_constraint_name!, name)
+              # rubocop:enable GitlabSecurity/PublicSend
+
+              column_name = helper_context.quote_column_name(column_name)
+              definition = "char_length(#{column_name}) <= #{limit}"
+
+              check_constraints << { name: name, definition: definition }
+            end
+
+            t.instance_eval(&block) unless block.nil?
+          end
+
+          next if check_constraints.empty?
+
+          constraint_clauses = check_constraints.map do |constraint|
+            "ADD CONSTRAINT #{quote_table_name(constraint[:name])} CHECK (#{constraint[:definition]})"
+          end
+
+          execute(<<~SQL)
+            ALTER TABLE #{quote_table_name(table_name)}
+            #{constraint_clauses.join(",\n")}
+          SQL
+        end
+      end
+
       # Creates a new index, concurrently
       #
       # Example:

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::MigrationHelpers do
+  include Database::TableSchemaHelpers
+
   let(:model) do
     ActiveRecord::Migration.new.extend(described_class)
   end
@@ -96,6 +98,131 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#create_table_with_constraints' do
+    let(:table_name) { :test_table }
+    let(:column_attributes) do
+      [
+        { name: 'id',         sql_type: 'bigint',                   null: false, default: nil    },
+        { name: 'created_at', sql_type: 'timestamp with time zone', null: false, default: nil    },
+        { name: 'updated_at', sql_type: 'timestamp with time zone', null: false, default: nil    },
+        { name: 'some_id',    sql_type: 'integer',                  null: false, default: nil    },
+        { name: 'active',     sql_type: 'boolean',                  null: false, default: 'true' },
+        { name: 'name',       sql_type: 'text',                     null: true,  default: nil    }
+      ]
+    end
+
+    before do
+      allow(model).to receive(:transaction_open?).and_return(true)
+    end
+
+    context 'when no check constraints are defined' do
+      it 'creates the table as expected' do
+        model.create_table_with_constraints table_name do |t|
+          t.timestamps_with_timezone
+          t.integer :some_id, null: false
+          t.boolean :active, null: false, default: true
+          t.text :name
+        end
+
+        expect_table_columns_to_match(column_attributes, table_name)
+      end
+    end
+
+    context 'when check constraints are defined' do
+      context 'when the text_limit is explicity named' do
+        it 'creates the table as expected' do
+          model.create_table_with_constraints table_name do |t|
+            t.timestamps_with_timezone
+            t.integer :some_id, null: false
+            t.boolean :active, null: false, default: true
+            t.text :name
+
+            t.text_limit :name, 255, name: 'check_name_length'
+            t.check_constraint :some_id_is_positive, 'some_id > 0'
+          end
+
+          expect_table_columns_to_match(column_attributes, table_name)
+
+          expect_check_constraint(table_name, 'check_name_length', 'char_length(name) <= 255')
+          expect_check_constraint(table_name, 'some_id_is_positive', 'some_id > 0')
+        end
+      end
+
+      context 'when the text_limit is not named' do
+        it 'creates the table as expected, naming the text limit' do
+          model.create_table_with_constraints table_name do |t|
+            t.timestamps_with_timezone
+            t.integer :some_id, null: false
+            t.boolean :active, null: false, default: true
+            t.text :name
+
+            t.text_limit :name, 255
+            t.check_constraint :some_id_is_positive, 'some_id > 0'
+          end
+
+          expect_table_columns_to_match(column_attributes, table_name)
+
+          expect_check_constraint(table_name, 'check_cda6f69506', 'char_length(name) <= 255')
+          expect_check_constraint(table_name, 'some_id_is_positive', 'some_id > 0')
+        end
+      end
+
+      it 'runs the change within a with_lock_retries' do
+        expect(model).to receive(:with_lock_retries).ordered.and_yield
+        expect(model).to receive(:create_table).ordered.and_call_original
+        expect(model).to receive(:execute).with(<<~SQL).ordered
+          ALTER TABLE "#{table_name}"\nADD CONSTRAINT "check_cda6f69506" CHECK (char_length("name") <= 255)
+        SQL
+
+        model.create_table_with_constraints table_name do |t|
+          t.text :name
+          t.text_limit :name, 255
+        end
+      end
+
+      context 'when constraints are given invalid names' do
+        let(:expected_max_length) { described_class::MAX_IDENTIFIER_NAME_LENGTH }
+        let(:expected_error_message) { "The maximum allowed constraint name is #{expected_max_length} characters" }
+
+        context 'when the explicit text limit name is not valid' do
+          it 'raises an error' do
+            too_long_length = expected_max_length + 1
+
+            expect do
+              model.create_table_with_constraints table_name do |t|
+                t.timestamps_with_timezone
+                t.integer :some_id, null: false
+                t.boolean :active, null: false, default: true
+                t.text :name
+
+                t.text_limit :name, 255, name: ('a' * too_long_length)
+                t.check_constraint :some_id_is_positive, 'some_id > 0'
+              end
+            end.to raise_error(expected_error_message)
+          end
+        end
+
+        context 'when a check constraint name is not valid' do
+          it 'raises an error' do
+            too_long_length = expected_max_length + 1
+
+            expect do
+              model.create_table_with_constraints table_name do |t|
+                t.timestamps_with_timezone
+                t.integer :some_id, null: false
+                t.boolean :active, null: false, default: true
+                t.text :name
+
+                t.text_limit :name, 255
+                t.check_constraint ('a' * too_long_length), 'some_id > 0'
+              end
+            end.to raise_error(expected_error_message)
+          end
+        end
+      end
+    end
+  end
+
   describe '#add_concurrent_index' do
     context 'outside a transaction' do
       before do

spec/support/helpers/database/table_schema_helpers.rb

@@ -17,6 +17,24 @@ module Database
       expect(table_oid(replacement_table)).to be_nil
     end
 
+    def expect_table_columns_to_match(expected_column_attributes, table_name)
+      expect(connection.table_exists?(table_name)).to eq(true)
+
+      actual_columns = connection.columns(table_name)
+      expect(actual_columns.size).to eq(column_attributes.size)
+
+      column_attributes.each_with_index do |attributes, i|
+        actual_column = actual_columns[i]
+
+        attributes.each do |name, value|
+          actual_value = actual_column.public_send(name)
+          message = "expected #{actual_column.name}.#{name} to be #{value}, but got #{actual_value}"
+
+          expect(actual_value).to eq(value), message
+        end
+      end
+    end
+
     def expect_index_to_exist(name, schema: nil)
       expect(index_exists_by_name(name, schema: schema)).to eq(true)
     end
@@ -25,6 +43,10 @@ module Database
       expect(index_exists_by_name(name, schema: schema)).to be_nil
     end
 
+    def expect_check_constraint(table_name, name, definition, schema: nil)
+      expect(check_constraint_definition(table_name, name, schema: schema)).to eq("CHECK ((#{definition}))")
+    end
+
     def expect_primary_keys_after_tables(tables, schema: nil)
       tables.each do |table|
         primary_key = primary_key_constraint_name(table, schema: schema)
@@ -110,5 +132,18 @@ module Database
           AND n.nspname = #{schema}
       SQL
     end
+
+    def check_constraint_definition(table_name, constraint_name, schema: nil)
+      table_name = schema ? "#{schema}.#{table_name}" : table_name
+
+      connection.select_value(<<~SQL)
+        SELECT
+          pg_get_constraintdef(oid) AS constraint_definition
+        FROM pg_catalog.pg_constraint
+        WHERE pg_constraint.conrelid = '#{table_name}'::regclass
+          AND pg_constraint.contype = 'c'
+          AND pg_constraint.conname = '#{constraint_name}'
+      SQL
+    end
   end
 end