diff --git a/db/docs/batched_background_migrations/backfill_cluster_providers_aws_sharding_key.yml b/db/docs/batched_background_migrations/backfill_cluster_providers_aws_sharding_key.yml new file mode 100644 index 0000000000000000000000000000000000000000..0240f65f3606b36889da7ed7c940fa3f4bfff5a2 --- /dev/null +++ b/db/docs/batched_background_migrations/backfill_cluster_providers_aws_sharding_key.yml @@ -0,0 +1,8 @@ +--- +migration_job_name: BackfillClusterProvidersAwsShardingKey +description: Backfill sharding key columns for cluster_providers_aws +feature_category: deployment_management +introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/216884 +milestone: '18.8' +queued_migration_version: 20251217000016 +finalized_by: # version of the migration that finalized this BBM diff --git a/db/docs/cluster_providers_aws.yml b/db/docs/cluster_providers_aws.yml index 3a34cf072e7dcd00a2ad8cff61e6e6291014c4c0..d47ef0654ec8400c18a513fab79063ba00b73bbd 100644 --- a/db/docs/cluster_providers_aws.yml +++ b/db/docs/cluster_providers_aws.yml @@ -33,5 +33,5 @@ desired_sharding_key: table: clusters sharding_key: organization_id belongs_to: cluster -sharding_key_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/583002 +sharding_key_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/584431 table_size: small diff --git a/db/migrate/20251217000009_add_sharding_key_columns_to_cluster_providers_aws.rb b/db/migrate/20251217000009_add_sharding_key_columns_to_cluster_providers_aws.rb new file mode 100644 index 0000000000000000000000000000000000000000..680d5935bc4801c8782d8ba398742906b91e8941 --- /dev/null +++ b/db/migrate/20251217000009_add_sharding_key_columns_to_cluster_providers_aws.rb @@ -0,0 +1,19 @@ +# frozen_string_literal: true + +class AddShardingKeyColumnsToClusterProvidersAws < Gitlab::Database::Migration[2.3] + milestone '18.8' + + TABLE_NAME = 'cluster_providers_aws' + + def up + add_column TABLE_NAME, :organization_id, :bigint + add_column TABLE_NAME, :namespace_id, :bigint + add_column TABLE_NAME, :project_id, :bigint + end + + def down + remove_column TABLE_NAME, :project_id + remove_column TABLE_NAME, :namespace_id + remove_column TABLE_NAME, :organization_id + end +end diff --git a/db/post_migrate/20251217000010_add_sharding_key_trigger_on_cluster_providers_aws.rb b/db/post_migrate/20251217000010_add_sharding_key_trigger_on_cluster_providers_aws.rb new file mode 100644 index 0000000000000000000000000000000000000000..485ac8dbe26d2a30db10db09b864402617b2afd5 --- /dev/null +++ b/db/post_migrate/20251217000010_add_sharding_key_trigger_on_cluster_providers_aws.rb @@ -0,0 +1,39 @@ +# frozen_string_literal: true + +class AddShardingKeyTriggerOnClusterProvidersAws < Gitlab::Database::Migration[2.3] + milestone '18.8' + include Gitlab::Database::SchemaHelpers + + TABLE_NAME = 'cluster_providers_aws' + TRIGGER_FUNCTION_NAME = 'cluster_providers_aws_sharding_key' + TRIGGER_NAME = "trigger_#{TRIGGER_FUNCTION_NAME}" + + def up + execute(<<~SQL) + CREATE OR REPLACE FUNCTION #{TRIGGER_FUNCTION_NAME}() RETURNS TRIGGER AS $$ + BEGIN + IF num_nonnulls(NEW.organization_id, NEW.namespace_id, NEW.project_id) != 1 THEN + SELECT "organization_id", "group_id", "project_id" + INTO NEW."organization_id", NEW."namespace_id", NEW."project_id" + FROM "clusters" + WHERE "clusters"."id" = NEW."cluster_id"; + END IF; + + RETURN NEW; + END + $$ LANGUAGE PLPGSQL + SQL + + create_trigger( + TABLE_NAME, + TRIGGER_NAME, + TRIGGER_FUNCTION_NAME, + fires: 'BEFORE INSERT OR UPDATE' + ) + end + + def down + drop_trigger(TABLE_NAME, TRIGGER_NAME) + drop_function(TRIGGER_FUNCTION_NAME) + end +end diff --git a/db/post_migrate/20251217000011_add_concurrent_index_to_sharding_key_columns_on_cluster_providers_aws.rb b/db/post_migrate/20251217000011_add_concurrent_index_to_sharding_key_columns_on_cluster_providers_aws.rb new file mode 100644 index 0000000000000000000000000000000000000000..c34ea30256bf3b53f419dc8cd9e5f1284397f263 --- /dev/null +++ b/db/post_migrate/20251217000011_add_concurrent_index_to_sharding_key_columns_on_cluster_providers_aws.rb @@ -0,0 +1,24 @@ +# frozen_string_literal: true + +class AddConcurrentIndexToShardingKeyColumnsOnClusterProvidersAws < Gitlab::Database::Migration[2.3] + disable_ddl_transaction! + milestone '18.8' + + TABLE_NAME = 'cluster_providers_aws' + + ORGANIZATION_INDEX = 'idx_cluster_providers_aws_on_organization_id' + NAMESPACE_INDEX = 'idx_cluster_providers_aws_on_namespace_id' + PROJECT_INDEX = 'idx_cluster_providers_aws_on_project_id' + + def up + add_concurrent_index TABLE_NAME, :organization_id, name: ORGANIZATION_INDEX + add_concurrent_index TABLE_NAME, :namespace_id, name: NAMESPACE_INDEX + add_concurrent_index TABLE_NAME, :project_id, name: PROJECT_INDEX + end + + def down + remove_concurrent_index_by_name TABLE_NAME, PROJECT_INDEX + remove_concurrent_index_by_name TABLE_NAME, NAMESPACE_INDEX + remove_concurrent_index_by_name TABLE_NAME, ORGANIZATION_INDEX + end +end diff --git a/db/post_migrate/20251217000012_add_foreign_key_constraint_on_cluster_providers_aws_organization_id.rb b/db/post_migrate/20251217000012_add_foreign_key_constraint_on_cluster_providers_aws_organization_id.rb new file mode 100644 index 0000000000000000000000000000000000000000..9acf486ad37667390ef32c6817c44fe37e292d18 --- /dev/null +++ b/db/post_migrate/20251217000012_add_foreign_key_constraint_on_cluster_providers_aws_organization_id.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +class AddForeignKeyConstraintOnClusterProvidersAwsOrganizationId < Gitlab::Database::Migration[2.3] + disable_ddl_transaction! + milestone '18.8' + + TABLE_NAME = 'cluster_providers_aws' + + def up + add_concurrent_foreign_key TABLE_NAME, + :organizations, + column: :organization_id, + on_delete: :cascade, + validate: false + end + + def down + remove_foreign_key_if_exists TABLE_NAME, :organizations, column: :organization_id + end +end diff --git a/db/post_migrate/20251217000013_add_foreign_key_constraint_on_cluster_providers_aws_namespace_id.rb b/db/post_migrate/20251217000013_add_foreign_key_constraint_on_cluster_providers_aws_namespace_id.rb new file mode 100644 index 0000000000000000000000000000000000000000..26ec564ceb05ed6873d4f71224b4f224e02066cb --- /dev/null +++ b/db/post_migrate/20251217000013_add_foreign_key_constraint_on_cluster_providers_aws_namespace_id.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +class AddForeignKeyConstraintOnClusterProvidersAwsNamespaceId < Gitlab::Database::Migration[2.3] + disable_ddl_transaction! + milestone '18.8' + + TABLE_NAME = 'cluster_providers_aws' + + def up + add_concurrent_foreign_key TABLE_NAME, + :namespaces, + column: :namespace_id, + on_delete: :cascade, + validate: false + end + + def down + remove_foreign_key_if_exists TABLE_NAME, :namespaces, column: :namespace_id + end +end diff --git a/db/post_migrate/20251217000014_add_foreign_key_constraint_on_cluster_providers_aws_project_id.rb b/db/post_migrate/20251217000014_add_foreign_key_constraint_on_cluster_providers_aws_project_id.rb new file mode 100644 index 0000000000000000000000000000000000000000..c193625cd9dfc0533bf8a32c41f0c636ac3d57b2 --- /dev/null +++ b/db/post_migrate/20251217000014_add_foreign_key_constraint_on_cluster_providers_aws_project_id.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +class AddForeignKeyConstraintOnClusterProvidersAwsProjectId < Gitlab::Database::Migration[2.3] + disable_ddl_transaction! + milestone '18.8' + + TABLE_NAME = 'cluster_providers_aws' + + def up + add_concurrent_foreign_key TABLE_NAME, + :projects, + column: :project_id, + on_delete: :cascade, + validate: false + end + + def down + remove_foreign_key_if_exists TABLE_NAME, :projects, column: :project_id + end +end diff --git a/db/post_migrate/20251217000015_add_not_null_not_valid_constraint_on_cluster_providers_aws_sharding_key.rb b/db/post_migrate/20251217000015_add_not_null_not_valid_constraint_on_cluster_providers_aws_sharding_key.rb new file mode 100644 index 0000000000000000000000000000000000000000..8788a86c56d9274bcbb13c4102ab405807581907 --- /dev/null +++ b/db/post_migrate/20251217000015_add_not_null_not_valid_constraint_on_cluster_providers_aws_sharding_key.rb @@ -0,0 +1,23 @@ +# frozen_string_literal: true + +class AddNotNullNotValidConstraintOnClusterProvidersAwsShardingKey < Gitlab::Database::Migration[2.3] + disable_ddl_transaction! + milestone '18.8' + + TABLE_NAME = 'cluster_providers_aws' + + def up + add_multi_column_not_null_constraint( + TABLE_NAME, + :organization_id, :namespace_id, :project_id, + validate: false + ) + end + + def down + remove_multi_column_not_null_constraint( + TABLE_NAME, + :organization_id, :namespace_id, :project_id + ) + end +end diff --git a/db/post_migrate/20251217000016_queue_backfill_cluster_providers_aws_sharding_key.rb b/db/post_migrate/20251217000016_queue_backfill_cluster_providers_aws_sharding_key.rb new file mode 100644 index 0000000000000000000000000000000000000000..4f44e72c8501eebe37a7acafd06d4622e98a8764 --- /dev/null +++ b/db/post_migrate/20251217000016_queue_backfill_cluster_providers_aws_sharding_key.rb @@ -0,0 +1,24 @@ +# frozen_string_literal: true + +class QueueBackfillClusterProvidersAwsShardingKey < Gitlab::Database::Migration[2.3] + milestone '18.8' + restrict_gitlab_migration gitlab_schema: :gitlab_main_org + + MIGRATION = "BackfillClusterProvidersAwsShardingKey" + BATCH_SIZE = 1000 + SUB_BATCH_SIZE = 100 + + def up + queue_batched_background_migration( + MIGRATION, + :cluster_providers_aws, + :id, + batch_size: BATCH_SIZE, + sub_batch_size: SUB_BATCH_SIZE + ) + end + + def down + delete_batched_background_migration(MIGRATION, :cluster_providers_aws, :id, []) + end +end diff --git a/db/schema_migrations/20251217000009 b/db/schema_migrations/20251217000009 new file mode 100644 index 0000000000000000000000000000000000000000..b3d2266742b810bfde24bcaac8876a112dd710bc --- /dev/null +++ b/db/schema_migrations/20251217000009 @@ -0,0 +1 @@ +76a10f643dbea08de377ac5c28b789122de6a8d66b71827ddd5f19348ab106cc \ No newline at end of file diff --git a/db/schema_migrations/20251217000010 b/db/schema_migrations/20251217000010 new file mode 100644 index 0000000000000000000000000000000000000000..fab0b46d551275936dc85f289cfea299fd2b1323 --- /dev/null +++ b/db/schema_migrations/20251217000010 @@ -0,0 +1 @@ +deb607356052b234f6ee92584fdefff27cd0c7b0e719bfb9552da340dc012d7c \ No newline at end of file diff --git a/db/schema_migrations/20251217000011 b/db/schema_migrations/20251217000011 new file mode 100644 index 0000000000000000000000000000000000000000..7deba1332f78887b184453bdf0b6728a1284350b --- /dev/null +++ b/db/schema_migrations/20251217000011 @@ -0,0 +1 @@ +0b582f43f5789253114c2849dc12eb20a135adb58e52f73d365af8c093e88b95 \ No newline at end of file diff --git a/db/schema_migrations/20251217000012 b/db/schema_migrations/20251217000012 new file mode 100644 index 0000000000000000000000000000000000000000..66666148b326778b80b247df9f0ac166fb10dcfc --- /dev/null +++ b/db/schema_migrations/20251217000012 @@ -0,0 +1 @@ +0fc0e0883ee2c9762d6f676040f21eac494c5d4d1b444604f461e1ef4fc1c2af \ No newline at end of file diff --git a/db/schema_migrations/20251217000013 b/db/schema_migrations/20251217000013 new file mode 100644 index 0000000000000000000000000000000000000000..81135527f657fddb2a0e7f7cfc2b2121e9a010be --- /dev/null +++ b/db/schema_migrations/20251217000013 @@ -0,0 +1 @@ +c06444182bd65fc8ba103163ad86608e4ae2133eedb5a682a1c28a37b86c23f1 \ No newline at end of file diff --git a/db/schema_migrations/20251217000014 b/db/schema_migrations/20251217000014 new file mode 100644 index 0000000000000000000000000000000000000000..3de954ceb7417cfb413bef65205df425b69c14b5 --- /dev/null +++ b/db/schema_migrations/20251217000014 @@ -0,0 +1 @@ +7c74a6bbcf8be0018ffd2e5553e17237c5463c7c115a4ffb2d7364fa09ef4673 \ No newline at end of file diff --git a/db/schema_migrations/20251217000015 b/db/schema_migrations/20251217000015 new file mode 100644 index 0000000000000000000000000000000000000000..21591eb3dbc4560a6deba9ce2df172db5e4f8fd4 --- /dev/null +++ b/db/schema_migrations/20251217000015 @@ -0,0 +1 @@ +31a2287c171e11fbde9d066214eb24e25dd6fd81602089008e66b442895326fe \ No newline at end of file diff --git a/db/schema_migrations/20251217000016 b/db/schema_migrations/20251217000016 new file mode 100644 index 0000000000000000000000000000000000000000..95f8d9b4a8bbc98c46cc83a6b5e71f0cd658ab43 --- /dev/null +++ b/db/schema_migrations/20251217000016 @@ -0,0 +1 @@ +58cc2cc90f564ec08352e44a6dc7abdbca68b09ab0498e5d662044233daa4922 \ No newline at end of file diff --git a/db/structure.sql b/db/structure.sql index 3059ddcc7554d1720dd64d469ebf167c25e6521a..12e59c70d1232c55643fe6330224e86b0a6b0c0f 100644 --- a/db/structure.sql +++ b/db/structure.sql @@ -234,6 +234,21 @@ BEGIN END; $$; +CREATE FUNCTION cluster_providers_aws_sharding_key() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + IF num_nonnulls(NEW.organization_id, NEW.namespace_id, NEW.project_id) != 1 THEN + SELECT "organization_id", "group_id", "project_id" + INTO NEW."organization_id", NEW."namespace_id", NEW."project_id" + FROM "clusters" + WHERE "clusters"."id" = NEW."cluster_id"; + END IF; + + RETURN NEW; +END +$$; + CREATE FUNCTION custom_dashboard_search_vector_update() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -15567,6 +15582,9 @@ CREATE TABLE cluster_providers_aws ( session_token text, status_reason text, kubernetes_version text DEFAULT '1.14'::text NOT NULL, + organization_id bigint, + namespace_id bigint, + project_id bigint, CONSTRAINT check_f1f42cd85e CHECK ((char_length(kubernetes_version) <= 30)) ); @@ -34889,6 +34907,9 @@ ALTER TABLE suggestions ALTER TABLE note_diff_files ADD CONSTRAINT check_ebb23d73d7 CHECK ((namespace_id IS NOT NULL)) NOT VALID; +ALTER TABLE cluster_providers_aws + ADD CONSTRAINT check_fca4a4cb60 CHECK ((num_nonnulls(namespace_id, organization_id, project_id) = 1)) NOT VALID; + ALTER TABLE vulnerability_statistics ADD CONSTRAINT check_vulnerability_statistics_traversal_ids_not_empty CHECK ((cardinality(traversal_ids) > 0)) NOT VALID; @@ -39481,6 +39502,12 @@ CREATE UNIQUE INDEX idx_ci_runner_taggings_proj_type_on_tag_id_runner_id_and_typ CREATE INDEX idx_ci_running_builds_on_runner_type_and_owner_xid_and_id ON ci_running_builds USING btree (runner_type, runner_owner_namespace_xid, runner_id); +CREATE INDEX idx_cluster_providers_aws_on_namespace_id ON cluster_providers_aws USING btree (namespace_id); + +CREATE INDEX idx_cluster_providers_aws_on_organization_id ON cluster_providers_aws USING btree (organization_id); + +CREATE INDEX idx_cluster_providers_aws_on_project_id ON cluster_providers_aws USING btree (project_id); + CREATE INDEX idx_compliance_requirements_controls_on_namespace_id ON compliance_requirements_controls USING btree (namespace_id); CREATE INDEX idx_compliance_requirements_controls_on_requirement_id ON compliance_requirements_controls USING btree (compliance_requirement_id); @@ -49259,6 +49286,8 @@ CREATE TRIGGER trigger_cfbec3f07e2b BEFORE INSERT OR UPDATE ON deployment_merge_ CREATE TRIGGER trigger_cleanup_pipeline_iid_after_delete AFTER DELETE ON p_ci_pipelines FOR EACH ROW EXECUTE FUNCTION cleanup_pipeline_iid_after_delete(); +CREATE TRIGGER trigger_cluster_providers_aws_sharding_key BEFORE INSERT OR UPDATE ON cluster_providers_aws FOR EACH ROW EXECUTE FUNCTION cluster_providers_aws_sharding_key(); + CREATE TRIGGER trigger_d32ff9d5c63d BEFORE INSERT OR UPDATE ON bulk_import_export_upload_uploads FOR EACH ROW EXECUTE FUNCTION trigger_d32ff9d5c63d(); CREATE TRIGGER trigger_d4487a75bd44 BEFORE INSERT OR UPDATE ON terraform_state_versions FOR EACH ROW EXECUTE FUNCTION trigger_d4487a75bd44(); @@ -49821,6 +49850,9 @@ ALTER TABLE ONLY security_pipeline_execution_project_schedules ALTER TABLE p_ci_build_trace_metadata ADD CONSTRAINT fk_21d25cac1a_p FOREIGN KEY (partition_id, trace_artifact_id) REFERENCES p_ci_job_artifacts(partition_id, id) ON UPDATE CASCADE ON DELETE CASCADE; +ALTER TABLE ONLY cluster_providers_aws + ADD CONSTRAINT fk_22b9b8f491 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE NOT VALID; + ALTER TABLE ONLY users_star_projects ADD CONSTRAINT fk_22cd27ddfc FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; @@ -50868,6 +50900,9 @@ ALTER TABLE ONLY merge_request_diff_commit_users ALTER TABLE ONLY packages_pypi_metadata ADD CONSTRAINT fk_884056a10f FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; +ALTER TABLE ONLY cluster_providers_aws + ADD CONSTRAINT fk_88769dd4f2 FOREIGN KEY (namespace_id) REFERENCES namespaces(id) ON DELETE CASCADE NOT VALID; + ALTER TABLE ONLY approval_group_rules_users ADD CONSTRAINT fk_888a0df3b7 FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; @@ -51927,6 +51962,9 @@ ALTER TABLE ONLY epic_user_mentions ALTER TABLE ONLY board_user_preferences ADD CONSTRAINT fk_f1c3e9b710 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE; +ALTER TABLE ONLY cluster_providers_aws + ADD CONSTRAINT fk_f1c5a1b10f FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE NOT VALID; + ALTER TABLE ONLY observability_metrics_issues_connections ADD CONSTRAINT fk_f218d84a14 FOREIGN KEY (namespace_id) REFERENCES namespaces(id) ON DELETE CASCADE; diff --git a/lib/gitlab/background_migration/backfill_cluster_providers_aws_sharding_key.rb b/lib/gitlab/background_migration/backfill_cluster_providers_aws_sharding_key.rb new file mode 100644 index 0000000000000000000000000000000000000000..f5c054adf88f4e2151e0147bc82d850c1b81054d --- /dev/null +++ b/lib/gitlab/background_migration/backfill_cluster_providers_aws_sharding_key.rb @@ -0,0 +1,33 @@ +# frozen_string_literal: true + +module Gitlab + module BackgroundMigration + class BackfillClusterProvidersAwsShardingKey < BatchedMigrationJob + operation_name :backfill_cluster_providers_aws_sharding_key + feature_category :deployment_management + + def perform + each_sub_batch do |sub_batch| + start_id = sub_batch.minimum(:id) + end_id = sub_batch.maximum(:id) + + connection.execute(<<~SQL) + UPDATE cluster_providers_aws + SET + organization_id = clusters.organization_id, + namespace_id = clusters.group_id, + project_id = clusters.project_id + FROM clusters + WHERE cluster_providers_aws.cluster_id = clusters.id + AND cluster_providers_aws.id BETWEEN #{start_id} AND #{end_id} + AND num_nonnulls( + cluster_providers_aws.organization_id, + cluster_providers_aws.namespace_id, + cluster_providers_aws.project_id + ) != 1; + SQL + end + end + end + end +end diff --git a/spec/lib/gitlab/background_migration/backfill_cluster_providers_aws_sharding_key_spec.rb b/spec/lib/gitlab/background_migration/backfill_cluster_providers_aws_sharding_key_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..de34975e25f537d529645f6b01465c3a9103ab02 --- /dev/null +++ b/spec/lib/gitlab/background_migration/backfill_cluster_providers_aws_sharding_key_spec.rb @@ -0,0 +1,171 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::BackgroundMigration::BackfillClusterProvidersAwsShardingKey, feature_category: :deployment_management do + let(:connection) { ApplicationRecord.connection } + + let(:organizations) { table(:organizations) } + let(:namespaces) { table(:namespaces) } + let(:projects) { table(:projects) } + let(:clusters) { table(:clusters) } + let(:cluster_providers_aws) { table(:cluster_providers_aws) } + + let(:function_name) { 'cluster_providers_aws_sharding_key' } + let(:trigger_name) { "trigger_#{function_name}" } + let(:constraint_name) { 'check_fca4a4cb60' } + + let(:organization) { organizations.create!(name: 'name', path: 'path') } + let(:namespace) { namespaces.create!(name: 'name', path: 'path', organization_id: organization.id) } + let(:project) do + projects.create!(namespace_id: namespace.id, project_namespace_id: namespace.id, organization_id: organization.id) + end + + let(:cluster_belonging_to_organization) do + clusters.create!( + name: 'test-cluster', + cluster_type: 1, + organization_id: organization.id, + group_id: nil, + project_id: nil + ) + end + + let(:provider_without_sharding_key) do + drop_constraint_and_trigger + record = cluster_providers_aws.create!( + cluster_id: cluster_belonging_to_organization.id, + num_nodes: 3, + status: 0, + key_name: 'test-key', + role_arn: 'arn:aws:iam::123456789012:role/test-role', + region: 'us-east-1', + vpc_id: 'vpc-12345678', + subnet_ids: ['subnet-12345678'], + security_group_id: 'sg-12345678', + instance_type: 't2.medium', + organization_id: nil, + namespace_id: nil, + project_id: nil + ) + add_constraint_and_trigger + record + end + + let(:provider_with_partial_sharding_key) do + cluster_providers_aws.create!( + cluster_id: cluster_belonging_to_organization.id, + num_nodes: 3, + status: 0, + key_name: 'test-key', + role_arn: 'arn:aws:iam::123456789012:role/test-role', + region: 'us-east-1', + vpc_id: 'vpc-12345678', + subnet_ids: ['subnet-12345678'], + security_group_id: 'sg-12345678', + instance_type: 't2.medium', + organization_id: organization.id, + namespace_id: nil, + project_id: nil + ) + end + + let(:provider_with_complete_sharding_key) do + drop_constraint_and_trigger + record = cluster_providers_aws.create!( + cluster_id: cluster_belonging_to_organization.id, + num_nodes: 3, + status: 0, + key_name: 'test-key', + role_arn: 'arn:aws:iam::123456789012:role/test-role', + region: 'us-east-1', + vpc_id: 'vpc-12345678', + subnet_ids: ['subnet-12345678'], + security_group_id: 'sg-12345678', + instance_type: 't2.medium', + organization_id: organization.id, + namespace_id: namespace.id, + project_id: project.id + ) + add_constraint_and_trigger + record + end + + subject(:migration) do + described_class.new( + start_id: cluster_providers_aws.minimum(:id), + end_id: cluster_providers_aws.maximum(:id), + batch_table: :cluster_providers_aws, + batch_column: :id, + sub_batch_size: 1, + pause_ms: 0, + connection: ApplicationRecord.connection + ) + end + + describe "#perform" do + it 'backfills sharding key for records that do not have it' do + provider_without_sharding_key.reload + expect(provider_without_sharding_key.organization_id).to be_nil + expect(provider_without_sharding_key.namespace_id).to be_nil + expect(provider_without_sharding_key.project_id).to be_nil + + migration.perform + + provider_without_sharding_key.reload + expect(provider_without_sharding_key.organization_id).to eq(organization.id) + expect(provider_without_sharding_key.namespace_id).to be_nil + expect(provider_without_sharding_key.project_id).to be_nil + end + + it 'does not modify records that already have complete sharding key' do + provider_with_complete_sharding_key.reload + expect(provider_with_complete_sharding_key.organization_id).to eq(organization.id) + expect(provider_with_complete_sharding_key.namespace_id).to eq(namespace.id) + expect(provider_with_complete_sharding_key.project_id).to eq(project.id) + + migration.perform + + provider_with_complete_sharding_key.reload + expect(provider_with_complete_sharding_key.organization_id).to eq(organization.id) + expect(provider_with_complete_sharding_key.namespace_id).to be_nil + expect(provider_with_complete_sharding_key.project_id).to be_nil + end + + it 'backfills sharding key for records with partial sharding key' do + provider_with_partial_sharding_key.reload + expect(provider_with_partial_sharding_key.organization_id).to eq(organization.id) + expect(provider_with_partial_sharding_key.namespace_id).to be_nil + expect(provider_with_partial_sharding_key.project_id).to be_nil + + migration.perform + + provider_with_partial_sharding_key.reload + expect(provider_with_partial_sharding_key.organization_id).to eq(organization.id) + expect(provider_with_partial_sharding_key.namespace_id).to be_nil + expect(provider_with_partial_sharding_key.project_id).to be_nil + end + end + + private + + def drop_constraint_and_trigger + connection.execute( + <<~SQL + DROP TRIGGER IF EXISTS #{trigger_name} ON cluster_providers_aws; + + ALTER TABLE cluster_providers_aws DROP CONSTRAINT IF EXISTS #{constraint_name}; + SQL + ) + end + + def add_constraint_and_trigger + connection.execute( + <<~SQL + ALTER TABLE cluster_providers_aws ADD CONSTRAINT #{constraint_name} CHECK ((num_nonnulls(namespace_id, organization_id, project_id) = 1)) NOT VALID; + + CREATE TRIGGER #{trigger_name} BEFORE INSERT OR UPDATE ON cluster_providers_aws FOR EACH ROW EXECUTE FUNCTION #{function_name}(); + SQL + ) + end +end diff --git a/spec/lib/gitlab/organizations/sharding_key_spec.rb b/spec/lib/gitlab/organizations/sharding_key_spec.rb index 5499b288fe21645c332b0e72e07d5b9304005c30..980082610ec6fbcbfc653c38735b91e16f01eb79 100644 --- a/spec/lib/gitlab/organizations/sharding_key_spec.rb +++ b/spec/lib/gitlab/organizations/sharding_key_spec.rb @@ -286,6 +286,7 @@ "ci_runner_machines" => "https://gitlab.com/gitlab-org/gitlab/-/issues/525293", "instance_type_ci_runner_machines" => "https://gitlab.com/gitlab-org/gitlab/-/issues/525293", "clusters" => "https://gitlab.com/gitlab-org/gitlab/-/issues/553452", + "cluster_providers_aws" => "https://gitlab.com/gitlab-org/gitlab/-/issues/584431", "ci_runners" => "https://gitlab.com/gitlab-org/gitlab/-/issues/525293", "instance_type_ci_runners" => "https://gitlab.com/gitlab-org/gitlab/-/issues/525293", "ci_runner_taggings" => "https://gitlab.com/gitlab-org/gitlab/-/issues/525293", diff --git a/spec/migrations/queue_backfill_cluster_providers_aws_sharding_key_spec.rb b/spec/migrations/queue_backfill_cluster_providers_aws_sharding_key_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..3bcf7483aec3c2a9a0bad1eb56dd28e5ea5455eb --- /dev/null +++ b/spec/migrations/queue_backfill_cluster_providers_aws_sharding_key_spec.rb @@ -0,0 +1,35 @@ +# frozen_string_literal: true + +require 'spec_helper' +require_migration! + +RSpec.describe QueueBackfillClusterProvidersAwsShardingKey, feature_category: :deployment_management do + let(:migration) { described_class.new } + + describe '#up' do + it 'queues the batched background migration' do + expect(migration).to receive(:queue_batched_background_migration).with( + 'BackfillClusterProvidersAwsShardingKey', + :cluster_providers_aws, + :id, + batch_size: 1000, + sub_batch_size: 100 + ) + + migration.up + end + end + + describe '#down' do + it 'deletes the batched background migration' do + expect(migration).to receive(:delete_batched_background_migration).with( + 'BackfillClusterProvidersAwsShardingKey', + :cluster_providers_aws, + :id, + [] + ) + + migration.down + end + end +end