reifydb_catalog/store/primary_key/
find.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::{ColumnDef, PrimaryKeyDef, PrimaryKeyKey, QueryTransaction, SourceId, TableId, ViewId},
6	return_internal_error,
7};
8
9use crate::{
10	CatalogStore,
11	store::primary_key::layout::{primary_key, primary_key::deserialize_column_ids},
12};
13
14impl CatalogStore {
15	pub async fn find_primary_key(
16		rx: &mut impl QueryTransaction,
17		source: impl Into<SourceId>,
18	) -> crate::Result<Option<PrimaryKeyDef>> {
19		let source_id = source.into();
20
21		// Get the primary key ID for the table or view
22		// Virtual tables and ring buffers don't have primary keys
23		// stored separately
24		let primary_key_id = match source_id {
25			SourceId::Table(table_id) => match Self::get_table_pk_id(rx, table_id).await? {
26				Some(pk_id) => pk_id,
27				None => return Ok(None),
28			},
29			SourceId::View(view_id) => match Self::get_view_pk_id(rx, view_id).await? {
30				Some(pk_id) => pk_id,
31				None => return Ok(None),
32			},
33			SourceId::Flow(_) => {
34				// Flows don't have primary keys
35				return Ok(None);
36			}
37			SourceId::TableVirtual(_) => {
38				// Virtual tables don't have primary keys
39				return Ok(None);
40			}
41			SourceId::RingBuffer(ringbuffer_id) => {
42				match Self::get_ringbuffer_pk_id(rx, ringbuffer_id).await? {
43					Some(pk_id) => pk_id,
44					None => return Ok(None),
45				}
46			}
47			SourceId::Dictionary(_) => {
48				// Dictionaries don't have traditional primary keys
49				return Ok(None);
50			}
51		};
52
53		// Fetch the primary key details
54		let primary_key_multi = match rx.get(&PrimaryKeyKey::encoded(primary_key_id)).await? {
55			Some(multi) => multi,
56			None => return_internal_error!(format!(
57				"Primary key with ID {:?} referenced but not found",
58				primary_key_id
59			)),
60		};
61
62		// Deserialize column IDs
63		let column_ids_blob = primary_key::LAYOUT.get_blob(&primary_key_multi.values, primary_key::COLUMN_IDS);
64		let column_ids = deserialize_column_ids(&column_ids_blob);
65
66		// Fetch full ColumnDef for each column ID
67		let mut columns = Vec::new();
68		for column_id in column_ids {
69			let column_def = Self::get_column(rx, column_id).await?;
70			columns.push(ColumnDef {
71				id: column_def.id,
72				name: column_def.name,
73				constraint: column_def.constraint,
74				policies: column_def.policies,
75				index: column_def.index,
76				auto_increment: column_def.auto_increment,
77				dictionary_id: None,
78			});
79		}
80
81		Ok(Some(PrimaryKeyDef {
82			id: primary_key_id,
83			columns,
84		}))
85	}
86
87	#[inline]
88	pub async fn find_table_primary_key(
89		rx: &mut impl QueryTransaction,
90		table_id: TableId,
91	) -> crate::Result<Option<PrimaryKeyDef>> {
92		Self::find_primary_key(rx, table_id).await
93	}
94
95	#[inline]
96	pub async fn find_view_primary_key(
97		rx: &mut impl QueryTransaction,
98		view_id: ViewId,
99	) -> crate::Result<Option<PrimaryKeyDef>> {
100		Self::find_primary_key(rx, view_id).await
101	}
102}