Merge branch '30453-add-migration-helpers-for-limiting-text-columns' into 'master'

Add migration helpers for managing check constraints and limits on text columns

See merge request gitlab-org/gitlab!29181

lib/gitlab/database/migration_helpers.rb

@@ -1178,8 +1178,147 @@ into similar problems in the future (e.g. when new tables are created).
         end
       end
 
+      # Returns the name for a check constraint
+      #
+      # type:
+      # - Any value, as long as it is unique
+      # - Constraint names are unique per table in Postgres, and, additionally,
+      #   we can have multiple check constraints over a column
+      #   So we use the (table, column, type) triplet as a unique name
+      # - e.g. we use 'max_length' when adding checks for text limits
+      #        or 'not_null' when adding a NOT NULL constraint
+      #
+      def check_constraint_name(table, column, type)
+        identifier = "#{table}_#{column}_check_#{type}"
+        # Check concurrent_foreign_key_name() for info on why we use a hash
+        hashed_identifier = Digest::SHA256.hexdigest(identifier).first(10)
+
+        "check_#{hashed_identifier}"
+      end
+
+      def check_constraint_exists?(table, constraint_name)
+        # Constraint names are unique per table in Postgres, not per schema
+        # Two tables can have constraints with the same name, so we filter by
+        # the table name in addition to using the constraint_name
+        check_sql = <<~SQL
+          SELECT COUNT(*)
+          FROM pg_constraint
+          JOIN pg_class ON pg_constraint.conrelid = pg_class.oid
+          WHERE pg_constraint.contype = 'c'
+          AND pg_constraint.conname = '#{constraint_name}'
+          AND pg_class.relname = '#{table}'
+        SQL
+
+        connection.select_value(check_sql).positive?
+      end
+
+      # Adds a check constraint to a table
+      #
+      # This method is the generic helper for adding any check constraint
+      # More specialized helpers may use it (e.g. add_text_limit or add_not_null)
+      #
+      # This method only requires minimal locking:
+      # - The constraint is added using NOT VALID
+      #   This allows us to add the check constraint without validating it
+      # - The check will be enforced for new data (inserts) coming in
+      # - If `validate: true` the constraint is also validated
+      #   Otherwise, validate_check_constraint() can be used at a later stage
+      # - Check comments on add_concurrent_foreign_key for more info
+      #
+      # table  - The table the constraint will be added to
+      # check  - The check clause to add
+      #          e.g. 'char_length(name) <= 5' or 'store IS NOT NULL'
+      # constraint_name - The name of the check constraint (otherwise auto-generated)
+      #                   Should be unique per table (not per column)
+      # validate - Whether to validate the constraint in this call
+      #
+      # rubocop:disable Gitlab/RailsLogger
+      def add_check_constraint(table, check, constraint_name, validate: true)
+        # Transactions would result in ALTER TABLE locks being held for the
+        # duration of the transaction, defeating the purpose of this method.
+        if transaction_open?
+          raise 'add_check_constraint can not be run inside a transaction'
+        end
+
+        if check_constraint_exists?(table, constraint_name)
+          warning_message = <<~MESSAGE
+            Check constraint was not created because it exists already
+            (this may be due to an aborted migration or similar)
+            table: #{table}, check: #{check}, constraint name: #{constraint_name}
+          MESSAGE
+
+          Rails.logger.warn warning_message
+        else
+          # Only add the constraint without validating it
+          # Even though it is fast, ADD CONSTRAINT requires an EXCLUSIVE lock
+          # Use with_lock_retries to make sure that this operation
+          # will not timeout on tables accessed by many processes
+          with_lock_retries do
+            execute <<-EOF.strip_heredoc
+            ALTER TABLE #{table}
+            ADD CONSTRAINT #{constraint_name}
+            CHECK ( #{check} )
+            NOT VALID;
+            EOF
+          end
+        end
+
+        if validate
+          validate_check_constraint(table, constraint_name)
+        end
+      end
+
+      def validate_check_constraint(table, constraint_name)
+        unless check_constraint_exists?(table, constraint_name)
+          raise missing_schema_object_message(table, "check constraint", constraint_name)
+        end
+
+        disable_statement_timeout do
+          # VALIDATE CONSTRAINT only requires a SHARE UPDATE EXCLUSIVE LOCK
+          # It only conflicts with other validations and creating indexes
+          execute("ALTER TABLE #{table} VALIDATE CONSTRAINT #{constraint_name};")
+        end
+      end
+
+      def remove_check_constraint(table, constraint_name)
+        # DROP CONSTRAINT requires an EXCLUSIVE lock
+        # Use with_lock_retries to make sure that this will not timeout
+        with_lock_retries do
+          execute <<-EOF.strip_heredoc
+          ALTER TABLE #{table}
+          DROP CONSTRAINT IF EXISTS #{constraint_name}
+          EOF
+        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(
+          table,
+          "char_length(#{column}) <= #{limit}",
+          text_limit_name(table, column, name: constraint_name),
+          validate: validate
+        )
+      end
+
+      def validate_text_limit(table, column, constraint_name: nil)
+        validate_check_constraint(table, text_limit_name(table, column, name: constraint_name))
+      end
+
+      def remove_text_limit(table, column, constraint_name: nil)
+        remove_check_constraint(table, text_limit_name(table, column, name: constraint_name))
+      end
+
+      def check_text_limit_exists?(table, column, constraint_name: nil)
+        check_constraint_exists?(table, text_limit_name(table, column, name: constraint_name))
+      end
+
       private
 
+      def text_limit_name(table, column, name: nil)
+        name.presence || check_constraint_name(table, column, 'max_length')
+      end
+
       def missing_schema_object_message(table, type, name)
         <<~MESSAGE
           Could not find #{type} "#{name}" on table "#{table}" which was referenced during the migration.

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -2046,4 +2046,333 @@ describe Gitlab::Database::MigrationHelpers do
       model.bulk_migrate_in(10.minutes, [%w(Class hello world)])
     end
   end
+
+  describe '#check_constraint_name' do
+    it 'returns a valid constraint name' do
+      name = model.check_constraint_name(:this_is_a_very_long_table_name,
+                                         :with_a_very_long_column_name,
+                                         :with_a_very_long_type)
+
+      expect(name).to be_an_instance_of(String)
+      expect(name).to start_with('check_')
+      expect(name.length).to eq(16)
+    end
+  end
+
+  describe '#check_constraint_exists?' do
+    before do
+      ActiveRecord::Base.connection.execute(
+        'ALTER TABLE projects ADD CONSTRAINT check_1 CHECK (char_length(path) <= 5) NOT VALID'
+      )
+    end
+
+    after do
+      ActiveRecord::Base.connection.execute(
+        'ALTER TABLE projects DROP CONSTRAINT IF EXISTS check_1'
+      )
+    end
+
+    it 'returns true if a constraint exists' do
+      expect(model.check_constraint_exists?(:projects, 'check_1'))
+        .to be_truthy
+    end
+
+    it 'returns false if a constraint does not exist' do
+      expect(model.check_constraint_exists?(:projects, 'this_does_not_exist'))
+        .to be_falsy
+    end
+
+    it 'returns false if a constraint with the same name exists in another table' do
+      expect(model.check_constraint_exists?(:users, 'check_1'))
+        .to be_falsy
+    end
+  end
+
+  describe '#add_check_constraint' do
+    before do
+      allow(model).to receive(:check_constraint_exists?).and_return(false)
+    end
+
+    context 'inside a transaction' do
+      it 'raises an error' do
+        expect(model).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          model.add_check_constraint(
+            :test_table,
+            'name IS NOT NULL',
+            'check_name_not_null'
+          )
+        end.to raise_error(RuntimeError)
+      end
+    end
+
+    context 'outside a transaction' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      context 'when the constraint is already defined in the database' do
+        it 'does not create a constraint' do
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(true)
+
+          expect(model).not_to receive(:execute).with(/ADD CONSTRAINT/)
+
+          # setting validate: false to only focus on the ADD CONSTRAINT command
+          model.add_check_constraint(
+            :test_table,
+            'name IS NOT NULL',
+            'check_name_not_null',
+            validate: false
+          )
+        end
+      end
+
+      context 'when the constraint is not defined in the database' do
+        it 'creates the constraint' do
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
+
+          # setting validate: false to only focus on the ADD CONSTRAINT command
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null',
+            validate: false
+          )
+        end
+      end
+
+      context 'when validate is not provided' do
+        it 'performs validation' do
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(false).exactly(1)
+
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
+
+          # we need the check constraint to exist so that the validation proceeds
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(true).exactly(1)
+
+          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null'
+          )
+        end
+      end
+
+      context 'when validate is provided with a falsey value' do
+        it 'skips validation' do
+          expect(model).not_to receive(:disable_statement_timeout)
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT/)
+          expect(model).not_to receive(:execute).with(/VALIDATE CONSTRAINT/)
+
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null',
+            validate: false
+          )
+        end
+      end
+
+      context 'when validate is provided with a truthy value' do
+        it 'performs validation' do
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(false).exactly(1)
+
+          expect(model).to receive(:disable_statement_timeout).and_call_original
+          expect(model).to receive(:execute).with(/statement_timeout/)
+          expect(model).to receive(:with_lock_retries).and_call_original
+          expect(model).to receive(:execute).with(/ADD CONSTRAINT check_name_not_null/)
+
+          expect(model).to receive(:check_constraint_exists?)
+                       .with(:test_table, 'check_name_not_null')
+                       .and_return(true).exactly(1)
+
+          expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
+          expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+
+          model.add_check_constraint(
+            :test_table,
+            'char_length(name) <= 255',
+            'check_name_not_null',
+            validate: true
+          )
+        end
+      end
+    end
+  end
+
+  describe '#validate_check_constraint' do
+    context 'when the constraint does not exist' do
+      it 'raises an error' do
+        error_message = /Could not find check constraint "check_1" on table "test_table"/
+
+        expect(model).to receive(:check_constraint_exists?).and_return(false)
+
+        expect do
+          model.validate_check_constraint(:test_table, 'check_1')
+        end.to raise_error(RuntimeError, error_message)
+      end
+    end
+
+    context 'when the constraint exists' do
+      it 'performs validation' do
+        validate_sql = /ALTER TABLE test_table VALIDATE CONSTRAINT check_name/
+
+        expect(model).to receive(:check_constraint_exists?).and_return(true)
+        expect(model).to receive(:disable_statement_timeout).and_call_original
+        expect(model).to receive(:execute).with(/statement_timeout/)
+        expect(model).to receive(:execute).ordered.with(validate_sql)
+        expect(model).to receive(:execute).ordered.with(/RESET ALL/)
+
+        model.validate_check_constraint(:test_table, 'check_name')
+      end
+    end
+  end
+
+  describe '#remove_check_constraint' do
+    it 'removes the constraint' do
+      drop_sql = /ALTER TABLE test_table\s+DROP CONSTRAINT IF EXISTS check_name/
+
+      expect(model).to receive(:with_lock_retries).and_call_original
+      expect(model).to receive(:execute).with(drop_sql)
+
+      model.remove_check_constraint(:test_table, 'check_name')
+    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
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+        check = "char_length(name) <= 255"
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:add_check_constraint)
+                     .with(:test_table, check, constraint_name, validate: true)
+
+        model.add_text_limit(:test_table, :name, 255)
+      end
+    end
+
+    context 'when all parameters are provided' do
+      it 'calls add_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_limit'
+        check = "char_length(name) <= 255"
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:add_check_constraint)
+                     .with(:test_table, check, constraint_name, validate: false)
+
+        model.add_text_limit(
+          :test_table,
+          :name,
+          255,
+          constraint_name: constraint_name,
+          validate: false
+        )
+      end
+    end
+  end
+
+  describe '#validate_text_limit' do
+    context 'when constraint_name is not provided' do
+      it 'calls validate_check_constraint with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:validate_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.validate_text_limit(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls validate_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_limit'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:validate_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.validate_text_limit(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#remove_text_limit' do
+    context 'when constraint_name is not provided' do
+      it 'calls remove_check_constraint with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:remove_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.remove_text_limit(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls remove_check_constraint with the correct parameters' do
+        constraint_name = 'check_name_limit'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:remove_check_constraint)
+                     .with(:test_table, constraint_name)
+
+        model.remove_text_limit(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
+
+  describe '#check_text_limit_exists?' do
+    context 'when constraint_name is not provided' do
+      it 'calls check_constraint_exists? with an infered constraint name' do
+        constraint_name = model.check_constraint_name(:test_table,
+                                                      :name,
+                                                      'max_length')
+
+        expect(model).to receive(:check_constraint_name).and_call_original
+        expect(model).to receive(:check_constraint_exists?)
+                     .with(:test_table, constraint_name)
+
+        model.check_text_limit_exists?(:test_table, :name)
+      end
+    end
+
+    context 'when constraint_name is provided' do
+      it 'calls check_constraint_exists? with the correct parameters' do
+        constraint_name = 'check_name_limit'
+
+        expect(model).not_to receive(:check_constraint_name)
+        expect(model).to receive(:check_constraint_exists?)
+                     .with(:test_table, constraint_name)
+
+        model.check_text_limit_exists?(:test_table, :name, constraint_name: constraint_name)
+      end
+    end
+  end
 end