reifydb_catalog/store/table/
set_pk.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4use reifydb_core::{
5	interface::{CommandTransaction, PrimaryKeyId, TableId, TableKey},
6	return_internal_error,
7};
8
9use crate::{CatalogStore, store::table::layout::table};
10
11impl CatalogStore {
12	/// Set the primary key ID for a table
13	/// Returns an internal error if the table doesn't exist
14	pub async fn set_table_primary_key(
15		txn: &mut impl CommandTransaction,
16		table_id: TableId,
17		primary_key_id: PrimaryKeyId,
18	) -> crate::Result<()> {
19		let multi = match txn.get(&TableKey::encoded(table_id)).await? {
20			Some(v) => v,
21			None => return_internal_error!(format!(
22				"Table with ID {} not found when setting primary key. This indicates a critical catalog inconsistency.",
23				table_id.0
24			)),
25		};
26
27		let mut updated_row = multi.values.clone();
28		table::LAYOUT.set_u64(&mut updated_row, table::PRIMARY_KEY, primary_key_id.0);
29
30		txn.set(&TableKey::encoded(table_id), updated_row).await?;
31
32		Ok(())
33	}
34}
35
36#[cfg(test)]
37mod tests {
38	use reifydb_core::interface::{PrimaryKeyId, TableId};
39	use reifydb_engine::test_utils::create_test_command_transaction;
40
41	use crate::{CatalogStore, test_utils::ensure_test_table};
42
43	#[tokio::test]
44	async fn test_set_table_primary_key() {
45		let mut txn = create_test_command_transaction().await;
46		let table = ensure_test_table(&mut txn).await;
47
48		// Set primary key
49		CatalogStore::set_table_primary_key(&mut txn, table.id, PrimaryKeyId(42)).await.unwrap();
50
51		// The test succeeds if no error is thrown.
52		// In real usage, create_primary_key would create both the
53		// PrimaryKey record and update the table, and
54		// find_primary_key would find it.
55	}
56
57	#[tokio::test]
58	async fn test_set_table_primary_key_nonexistent() {
59		let mut txn = create_test_command_transaction().await;
60
61		// Try to set primary key on non-existent table
62		let result = CatalogStore::set_table_primary_key(&mut txn, TableId(999), PrimaryKeyId(1)).await;
63
64		assert!(result.is_err());
65		let err = result.unwrap_err();
66		assert!(err.to_string().contains("Table with ID 999 not found"));
67		assert!(err.to_string().contains("critical catalog inconsistency"));
68	}
69}
reifydb_catalog/store/table/set_pk.rs

reifydb_catalog/store/table/
set_pk.rs
