reifydb_catalog/store/table/
get_pk_id.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::interface::{PrimaryKeyId, QueryTransaction, TableId, TableKey};
5
6use crate::{CatalogStore, store::table::layout::table};
7
8impl CatalogStore {
9	/// Get the primary key ID for a table
10	/// Returns None if the table doesn't exist or has no primary key
11	pub async fn get_table_pk_id(
12		rx: &mut impl QueryTransaction,
13		table_id: TableId,
14	) -> crate::Result<Option<PrimaryKeyId>> {
15		let multi = match rx.get(&TableKey::encoded(table_id)).await? {
16			Some(v) => v,
17			None => return Ok(None),
18		};
19
20		let pk_id = table::LAYOUT.get_u64(&multi.values, table::PRIMARY_KEY);
21
22		if pk_id == 0 {
23			Ok(None)
24		} else {
25			Ok(Some(PrimaryKeyId(pk_id)))
26		}
27	}
28}
29
30#[cfg(test)]
31mod tests {
32	use reifydb_core::interface::{SourceId, TableId};
33	use reifydb_engine::test_utils::create_test_command_transaction;
34
35	use crate::{
36		CatalogStore,
37		column::{ColumnIndex, ColumnToCreate},
38		primary_key::PrimaryKeyToCreate,
39		test_utils::ensure_test_table,
40	};
41
42	#[tokio::test]
43	async fn test_get_table_pk_id_with_primary_key() {
44		let mut txn = create_test_command_transaction().await;
45		let table = ensure_test_table(&mut txn).await;
46
47		// Create a column
48		let col = CatalogStore::create_column(
49			&mut txn,
50			table.id,
51			ColumnToCreate {
52				fragment: None,
53				namespace_name: "test_namespace".to_string(),
54				table: table.id,
55				table_name: "test_table".to_string(),
56				column: "id".to_string(),
57				constraint: reifydb_type::TypeConstraint::unconstrained(reifydb_type::Type::Uint8),
58				if_not_exists: false,
59				policies: vec![],
60				index: ColumnIndex(0),
61				auto_increment: false,
62				dictionary_id: None,
63			},
64		)
65		.await
66		.unwrap();
67
68		// Create primary key
69		let pk_id = CatalogStore::create_primary_key(
70			&mut txn,
71			PrimaryKeyToCreate {
72				source: SourceId::Table(table.id),
73				column_ids: vec![col.id],
74			},
75		)
76		.await
77		.unwrap();
78
79		// Get the primary key ID
80		let retrieved_pk_id = CatalogStore::get_table_pk_id(&mut txn, table.id)
81			.await
82			.unwrap()
83			.expect("Primary key ID should exist");
84
85		assert_eq!(retrieved_pk_id, pk_id);
86	}
87
88	#[tokio::test]
89	async fn test_get_table_pk_id_without_primary_key() {
90		let mut txn = create_test_command_transaction().await;
91		let table = ensure_test_table(&mut txn).await;
92
93		// Get the primary key ID - should be None
94		let pk_id = CatalogStore::get_table_pk_id(&mut txn, table.id).await.unwrap();
95
96		assert!(pk_id.is_none());
97	}
98
99	#[tokio::test]
100	async fn test_get_table_pk_id_nonexistent_table() {
101		let mut txn = create_test_command_transaction().await;
102
103		// Get the primary key ID for non-existent table - should be
104		// None
105		let pk_id = CatalogStore::get_table_pk_id(&mut txn, TableId(999)).await.unwrap();
106
107		assert!(pk_id.is_none());
108	}
109}
reifydb_catalog/store/table/get_pk_id.rs

reifydb_catalog/store/table/
get_pk_id.rs
