diff --git a/db/post_migrate/20230119095022_add_fk_index_to_ci_builds_runner_session_on_partition_id_and_build_id.rb b/db/post_migrate/20230119095022_add_fk_index_to_ci_builds_runner_session_on_partition_id_and_build_id.rb
new file mode 100644
index 0000000000000000000000000000000000000000..9fee62e55f2af24d5f6be8905142affa5240f72e
--- /dev/null
+++ b/db/post_migrate/20230119095022_add_fk_index_to_ci_builds_runner_session_on_partition_id_and_build_id.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class AddFkIndexToCiBuildsRunnerSessionOnPartitionIdAndBuildId < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+
+  INDEX_NAME = :index_ci_builds_runner_session_on_partition_id_build_id
+  TABLE_NAME = :ci_builds_runner_session
+  COLUMNS = [:partition_id, :build_id]
+
+  def up
+    add_concurrent_index(TABLE_NAME, COLUMNS, name: INDEX_NAME, unique: true)
+  end
+
+  def down
+    remove_concurrent_index_by_name(TABLE_NAME, INDEX_NAME)
+  end
+end
diff --git a/db/post_migrate/20230119095023_add_fk_to_ci_builds_runner_session_on_partition_id_and_build_id.rb b/db/post_migrate/20230119095023_add_fk_to_ci_builds_runner_session_on_partition_id_and_build_id.rb
new file mode 100644
index 0000000000000000000000000000000000000000..e3aed238e0fa9b4438bc8daabbdd42cc0125fd58
--- /dev/null
+++ b/db/post_migrate/20230119095023_add_fk_to_ci_builds_runner_session_on_partition_id_and_build_id.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+class AddFkToCiBuildsRunnerSessionOnPartitionIdAndBuildId < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+
+  SOURCE_TABLE_NAME = :ci_builds_runner_session
+  TARGET_TABLE_NAME = :ci_builds
+  COLUMN = :build_id
+  TARGET_COLUMN = :id
+  FK_NAME = :fk_rails_70707857d3_p
+  PARTITION_COLUMN = :partition_id
+
+  def up
+    add_concurrent_foreign_key(
+      SOURCE_TABLE_NAME,
+      TARGET_TABLE_NAME,
+      column: [PARTITION_COLUMN, COLUMN],
+      target_column: [PARTITION_COLUMN, TARGET_COLUMN],
+      validate: false,
+      reverse_lock_order: true,
+      on_update: :cascade,
+      on_delete: :cascade,
+      name: FK_NAME
+    )
+  end
+
+  def down
+    with_lock_retries do
+      remove_foreign_key_if_exists(
+        SOURCE_TABLE_NAME,
+        TARGET_TABLE_NAME,
+        name: FK_NAME,
+        reverse_lock_order: true
+      )
+    end
+  end
+end
diff --git a/db/post_migrate/20230119095024_validate_fk_on_ci_builds_runner_session_partition_id_and_build_id.rb b/db/post_migrate/20230119095024_validate_fk_on_ci_builds_runner_session_partition_id_and_build_id.rb
new file mode 100644
index 0000000000000000000000000000000000000000..4c332a1db81c6b9177c0abbe3e1e9816be4c3049
--- /dev/null
+++ b/db/post_migrate/20230119095024_validate_fk_on_ci_builds_runner_session_partition_id_and_build_id.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class ValidateFkOnCiBuildsRunnerSessionPartitionIdAndBuildId < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+
+  TABLE_NAME = :ci_builds_runner_session
+  FK_NAME = :fk_rails_70707857d3_p
+  COLUMNS = [:partition_id, :build_id]
+
+  def up
+    validate_foreign_key(TABLE_NAME, COLUMNS, name: FK_NAME)
+  end
+
+  def down
+    # no-op
+  end
+end
diff --git a/db/post_migrate/20230119095025_remove_fk_to_ci_builds_ci_builds_runner_session_on_build_id.rb b/db/post_migrate/20230119095025_remove_fk_to_ci_builds_ci_builds_runner_session_on_build_id.rb
new file mode 100644
index 0000000000000000000000000000000000000000..0707114d3a5d0e07ff812c9d6c97b936b08367ae
--- /dev/null
+++ b/db/post_migrate/20230119095025_remove_fk_to_ci_builds_ci_builds_runner_session_on_build_id.rb
@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+class RemoveFkToCiBuildsCiBuildsRunnerSessionOnBuildId < Gitlab::Database::Migration[2.1]
+  disable_ddl_transaction!
+
+  SOURCE_TABLE_NAME = :ci_builds_runner_session
+  TARGET_TABLE_NAME = :ci_builds
+  COLUMN = :build_id
+  TARGET_COLUMN = :id
+  FK_NAME = :fk_rails_70707857d3
+
+  def up
+    with_lock_retries do
+      remove_foreign_key_if_exists(
+        SOURCE_TABLE_NAME,
+        TARGET_TABLE_NAME,
+        name: FK_NAME,
+        reverse_lock_order: true
+      )
+    end
+  end
+
+  def down
+    add_concurrent_foreign_key(
+      SOURCE_TABLE_NAME,
+      TARGET_TABLE_NAME,
+      column: COLUMN,
+      target_column: TARGET_COLUMN,
+      validate: true,
+      reverse_lock_order: true,
+      on_delete: :cascade,
+      name: FK_NAME
+    )
+  end
+end
diff --git a/db/schema_migrations/20230119095022 b/db/schema_migrations/20230119095022
new file mode 100644
index 0000000000000000000000000000000000000000..02c615d27f28f4fe4c0f9f45eec7bdca7f8cfe63
--- /dev/null
+++ b/db/schema_migrations/20230119095022
@@ -0,0 +1 @@
+938ad78781f488d0add8b1a85217f93ca41d7093015817d5bd5c8a853d3172b8
\ No newline at end of file
diff --git a/db/schema_migrations/20230119095023 b/db/schema_migrations/20230119095023
new file mode 100644
index 0000000000000000000000000000000000000000..ae882ff2ca7e1735f471a4748f27dae5778c006e
--- /dev/null
+++ b/db/schema_migrations/20230119095023
@@ -0,0 +1 @@
+f0b069e73170819e044496fb6eb02b43d0ee2697524659b50ab9b66dd3ec9792
\ No newline at end of file
diff --git a/db/schema_migrations/20230119095024 b/db/schema_migrations/20230119095024
new file mode 100644
index 0000000000000000000000000000000000000000..e7e4255c2331a3a776eaa73289664583edbede12
--- /dev/null
+++ b/db/schema_migrations/20230119095024
@@ -0,0 +1 @@
+be0b7fd198db3caa10e16606a861de6e708eadd225ad0360604a9f4c512a3d19
\ No newline at end of file
diff --git a/db/schema_migrations/20230119095025 b/db/schema_migrations/20230119095025
new file mode 100644
index 0000000000000000000000000000000000000000..f7eac79873c971e9697dca3dca60af1539174df7
--- /dev/null
+++ b/db/schema_migrations/20230119095025
@@ -0,0 +1 @@
+66239a6e2fdc476126ca247042078a1b939da15ff91c00cd0392ce664eadc9af
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index 04f0b7d18bb092ca8831d8a98a6d9d093bfb666a..d40a11351015e6502a67512bf77fee11338e5c7d 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -28925,6 +28925,8 @@ CREATE INDEX index_ci_builds_runner_id_running ON ci_builds USING btree (runner_
 
 CREATE UNIQUE INDEX index_ci_builds_runner_session_on_build_id ON ci_builds_runner_session USING btree (build_id);
 
+CREATE UNIQUE INDEX index_ci_builds_runner_session_on_partition_id_build_id ON ci_builds_runner_session USING btree (partition_id, build_id);
+
 CREATE INDEX index_ci_daily_build_group_report_results_on_group_id ON ci_daily_build_group_report_results USING btree (group_id);
 
 CREATE INDEX index_ci_daily_build_group_report_results_on_last_pipeline_id ON ci_daily_build_group_report_results USING btree (last_pipeline_id);
@@ -35008,7 +35010,7 @@ ALTER TABLE ONLY analytics_dashboards_pointers
     ADD CONSTRAINT fk_rails_7027b7eaa9 FOREIGN KEY (namespace_id) REFERENCES namespaces(id) ON DELETE CASCADE;
 
 ALTER TABLE ONLY ci_builds_runner_session
-    ADD CONSTRAINT fk_rails_70707857d3 FOREIGN KEY (build_id) REFERENCES ci_builds(id) ON DELETE CASCADE;
+    ADD CONSTRAINT fk_rails_70707857d3_p FOREIGN KEY (partition_id, build_id) REFERENCES ci_builds(partition_id, id) ON UPDATE CASCADE ON DELETE CASCADE;
 
 ALTER TABLE ONLY list_user_preferences
     ADD CONSTRAINT fk_rails_70b2ef5ce2 FOREIGN KEY (list_id) REFERENCES lists(id) ON DELETE CASCADE;
diff --git a/spec/db/schema_spec.rb b/spec/db/schema_spec.rb
index 4af901e4cf6975f65780dd4e4fa30567432fd3b3..9b839873fbb535b7059300061ecdcb0f58650b4f 100644
--- a/spec/db/schema_spec.rb
+++ b/spec/db/schema_spec.rb
@@ -38,7 +38,7 @@
     ci_build_trace_chunks: %w[partition_id],
     ci_build_trace_metadata: %w[partition_id],
     ci_builds: %w[erased_by_id trigger_request_id partition_id],
-    ci_builds_runner_session: %w[partition_id],
+    ci_builds_runner_session: %w[partition_id build_id],
     p_ci_builds_metadata: %w[partition_id],
     ci_job_artifacts: %w[partition_id],
     ci_job_variables: %w[partition_id],