1 files changed, 65 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/sharding_key_spec.rb b/spec/lib/gitlab/database/sharding_key_spec.rb
index ca01c7b8d36..c2e6f90f492 100644
--- a/spec/lib/gitlab/database/sharding_key_spec.rb
+++ b/spec/lib/gitlab/database/sharding_key_spec.rb
@@ -12,6 +12,15 @@ RSpec.describe 'new tables missing sharding_key', feature_category: :cell do
     ]
   end
 
+  # Specific tables can be temporarily exempt from this requirement. You must add an issue link in a comment next to
+  # the table name to remove this once a decision has been made.
+  let(:allowed_to_be_missing_not_null) do
+    [
+      'labels.project_id', # https://gitlab.com/gitlab-org/gitlab/-/issues/434356
+      'labels.group_id' # https://gitlab.com/gitlab-org/gitlab/-/issues/434356
+    ]
+  end
+
   let(:starting_from_milestone) { 16.6 }
 
   let(:allowed_sharding_key_referenced_tables) { %w[projects namespaces organizations] }
@@ -33,6 +42,26 @@ RSpec.describe 'new tables missing sharding_key', feature_category: :cell do
     end
   end
 
+  it 'ensures all sharding_key columns are not nullable or have a not null check constraint',
+    :aggregate_failures do
+    all_tables_to_sharding_key.each do |table_name, sharding_key|
+      sharding_key.each do |column_name, _|
+        not_nullable = not_nullable?(table_name, column_name)
+        has_null_check_constraint = has_null_check_constraint?(table_name, column_name)
+
+        if allowed_to_be_missing_not_null.include?("#{table_name}.#{column_name}")
+          expect(not_nullable || has_null_check_constraint).to eq(false),
+            "You must remove `#{table_name}.#{column_name}` from allowed_to_be_missing_not_null" \
+            "since it now has a valid constraint."
+        else
+          expect(not_nullable || has_null_check_constraint).to eq(true),
+            "Missing a not null constraint for `#{table_name}.#{column_name}` . " \
+            "All sharding keys must be not nullable or have a NOT NULL check constraint"
+        end
+      end
+    end
+  end
+
   it 'only allows `allowed_to_be_missing_sharding_key` to include tables that are missing a sharding_key',
     :aggregate_failures do
     allowed_to_be_missing_sharding_key.each do |exempted_table|
@@ -73,6 +102,42 @@ RSpec.describe 'new tables missing sharding_key', feature_category: :cell do
     end
   end
 
+  def not_nullable?(table_name, column_name)
+    sql = <<~SQL
+    SELECT 1
+    FROM information_schema.columns
+    WHERE table_schema = 'public' AND
+    table_name = '#{table_name}' AND
+    column_name = '#{column_name}' AND
+    is_nullable = 'NO'
+    SQL
+
+    result = ApplicationRecord.connection.execute(sql)
+
+    result.count > 0
+  end
+
+  def has_null_check_constraint?(table_name, column_name)
+    # This is a heuristic query to look for all check constraints on the table and see if any of them contain a clause
+    # column IS NOT NULL. This is to match tables that will have multiple sharding keys where either of them can be not
+    # null. Such cases may look like:
+    #    (project_id IS NOT NULL) OR (group_id IS NOT NULL)
+    # It's possible that this will sometimes incorrectly find a check constraint that isn't exactly as strict as we want
+    # but it should be pretty unlikely.
+    sql = <<~SQL
+    SELECT 1
+    FROM pg_constraint
+    INNER JOIN pg_class ON pg_constraint.conrelid = pg_class.oid
+    WHERE pg_class.relname = '#{table_name}'
+    AND contype = 'c'
+    AND pg_get_constraintdef(pg_constraint.oid) ILIKE '%#{column_name} IS NOT NULL%'
+    SQL
+
+    result = ApplicationRecord.connection.execute(sql)
+
+    result.count > 0
+  end
+
   def column_exists?(table_name, column_name)
     sql = <<~SQL
     SELECT 1