reifydb_catalog/store/primary_key/
find.rs1use 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 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 return Ok(None);
36 }
37 SourceId::TableVirtual(_) => {
38 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 return Ok(None);
50 }
51 };
52
53 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 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 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}