1use std::collections::{HashMap, HashSet};
16use std::fmt::Display;
17use std::sync::Arc;
18
19use common_base::bytes::Bytes;
20use serde::{Deserialize, Serialize};
21use snafu::OptionExt;
22use table::metadata::TableId;
23use table::table_name::TableName;
24
25use crate::error::{InvalidViewInfoSnafu, Result};
26use crate::key::txn_helper::TxnOpGetResponseSet;
27use crate::key::{
28 DeserializedValueWithBytes, MetadataKey, MetadataValue, VIEW_INFO_KEY_PATTERN,
29 VIEW_INFO_KEY_PREFIX,
30};
31use crate::kv_backend::KvBackendRef;
32use crate::kv_backend::txn::Txn;
33use crate::rpc::store::BatchGetRequest;
34
35#[derive(Debug, PartialEq)]
39pub struct ViewInfoKey {
40 view_id: TableId,
41}
42
43impl ViewInfoKey {
44 pub fn new(view_id: TableId) -> Self {
46 Self { view_id }
47 }
48}
49
50impl Display for ViewInfoKey {
51 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52 write!(f, "{}/{}", VIEW_INFO_KEY_PREFIX, self.view_id)
53 }
54}
55
56impl MetadataKey<'_, ViewInfoKey> for ViewInfoKey {
57 fn to_bytes(&self) -> Vec<u8> {
58 self.to_string().into_bytes()
59 }
60
61 fn from_bytes(bytes: &[u8]) -> Result<ViewInfoKey> {
62 let key = std::str::from_utf8(bytes).map_err(|e| {
63 InvalidViewInfoSnafu {
64 err_msg: format!(
65 "ViewInfoKey '{}' is not a valid UTF8 string: {e}",
66 String::from_utf8_lossy(bytes)
67 ),
68 }
69 .build()
70 })?;
71 let captures =
72 VIEW_INFO_KEY_PATTERN
73 .captures(key)
74 .with_context(|| InvalidViewInfoSnafu {
75 err_msg: format!("Invalid ViewInfoKey '{key}'"),
76 })?;
77 let view_id = captures[1].parse::<TableId>().unwrap();
79 Ok(ViewInfoKey { view_id })
80 }
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
85pub struct ViewInfoValue {
86 pub view_info: Bytes,
88 pub table_names: HashSet<TableName>,
90 pub columns: Vec<String>,
92 pub plan_columns: Vec<String>,
94 pub definition: String,
96 version: u64,
97}
98
99impl ViewInfoValue {
100 pub fn new(
101 view_info: Bytes,
102 table_names: HashSet<TableName>,
103 columns: Vec<String>,
104 plan_columns: Vec<String>,
105 definition: String,
106 ) -> Self {
107 Self {
108 view_info,
109 table_names,
110 columns,
111 plan_columns,
112 definition,
113 version: 0,
114 }
115 }
116
117 pub fn update(
120 &self,
121 new_view_info: Bytes,
122 table_names: HashSet<TableName>,
123 columns: Vec<String>,
124 plan_columns: Vec<String>,
125 definition: String,
126 ) -> Self {
127 Self {
128 view_info: new_view_info,
129 table_names,
130 columns,
131 plan_columns,
132 definition,
133 version: self.version + 1,
134 }
135 }
136}
137
138pub struct ViewInfoManager {
140 kv_backend: KvBackendRef,
141}
142
143pub type ViewInfoManagerRef = Arc<ViewInfoManager>;
144
145pub type ViewInfoValueDecodeResult = Result<Option<DeserializedValueWithBytes<ViewInfoValue>>>;
146
147impl ViewInfoManager {
148 pub fn new(kv_backend: KvBackendRef) -> Self {
149 Self { kv_backend }
150 }
151
152 pub(crate) fn build_create_txn(
154 &self,
155 view_id: TableId,
156 view_info_value: &ViewInfoValue,
157 ) -> Result<(
158 Txn,
159 impl FnOnce(&mut TxnOpGetResponseSet) -> ViewInfoValueDecodeResult,
160 )> {
161 let key = ViewInfoKey::new(view_id);
162 let raw_key = key.to_bytes();
163
164 let txn = Txn::put_if_not_exists(raw_key.clone(), view_info_value.try_as_raw_value()?);
165
166 Ok((
167 txn,
168 TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
169 ))
170 }
171
172 pub(crate) fn build_update_txn(
175 &self,
176 view_id: TableId,
177 current_view_info_value: &DeserializedValueWithBytes<ViewInfoValue>,
178 new_view_info_value: &ViewInfoValue,
179 ) -> Result<(
180 Txn,
181 impl FnOnce(&mut TxnOpGetResponseSet) -> ViewInfoValueDecodeResult,
182 )> {
183 let key = ViewInfoKey::new(view_id);
184 let raw_key = key.to_bytes();
185 let raw_value = current_view_info_value.get_raw_bytes();
186 let new_raw_value: Vec<u8> = new_view_info_value.try_as_raw_value()?;
187
188 let txn = Txn::compare_and_put(raw_key.clone(), raw_value, new_raw_value);
189
190 Ok((
191 txn,
192 TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
193 ))
194 }
195
196 pub async fn get(
198 &self,
199 view_id: TableId,
200 ) -> Result<Option<DeserializedValueWithBytes<ViewInfoValue>>> {
201 let key = ViewInfoKey::new(view_id);
202 let raw_key = key.to_bytes();
203 self.kv_backend
204 .get(&raw_key)
205 .await?
206 .map(|x| DeserializedValueWithBytes::from_inner_slice(&x.value))
207 .transpose()
208 }
209
210 pub async fn batch_get(&self, view_ids: &[TableId]) -> Result<HashMap<TableId, ViewInfoValue>> {
212 let lookup_table = view_ids
213 .iter()
214 .map(|id| (ViewInfoKey::new(*id).to_bytes(), id))
215 .collect::<HashMap<_, _>>();
216
217 let resp = self
218 .kv_backend
219 .batch_get(BatchGetRequest {
220 keys: lookup_table.keys().cloned().collect::<Vec<_>>(),
221 })
222 .await?;
223
224 let values = resp
225 .kvs
226 .iter()
227 .map(|kv| {
228 Ok((
229 **lookup_table.get(kv.key()).unwrap(),
231 ViewInfoValue::try_from_raw_value(&kv.value)?,
232 ))
233 })
234 .collect::<Result<HashMap<_, _>>>()?;
235
236 Ok(values)
237 }
238
239 pub async fn batch_get_raw(
241 &self,
242 view_ids: &[TableId],
243 ) -> Result<HashMap<TableId, DeserializedValueWithBytes<ViewInfoValue>>> {
244 let lookup_table = view_ids
245 .iter()
246 .map(|id| (ViewInfoKey::new(*id).to_bytes(), id))
247 .collect::<HashMap<_, _>>();
248
249 let resp = self
250 .kv_backend
251 .batch_get(BatchGetRequest {
252 keys: lookup_table.keys().cloned().collect::<Vec<_>>(),
253 })
254 .await?;
255
256 let values = resp
257 .kvs
258 .iter()
259 .map(|kv| {
260 Ok((
261 **lookup_table.get(kv.key()).unwrap(),
263 DeserializedValueWithBytes::from_inner_slice(&kv.value)?,
264 ))
265 })
266 .collect::<Result<HashMap<_, _>>>()?;
267
268 Ok(values)
269 }
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275
276 #[test]
277 fn test_key_serialization() {
278 let key = ViewInfoKey::new(42);
279 let raw_key = key.to_bytes();
280 assert_eq!(raw_key, b"__view_info/42");
281 }
282
283 #[test]
284 fn test_key_deserialization() {
285 let expected = ViewInfoKey::new(42);
286 let key = ViewInfoKey::from_bytes(b"__view_info/42").unwrap();
287 assert_eq!(key, expected);
288 }
289
290 #[test]
291 fn test_value_serialization() {
292 let table_names = {
293 let mut set = HashSet::new();
294 set.insert(TableName {
295 catalog_name: "greptime".to_string(),
296 schema_name: "public".to_string(),
297 table_name: "a_table".to_string(),
298 });
299 set.insert(TableName {
300 catalog_name: "greptime".to_string(),
301 schema_name: "public".to_string(),
302 table_name: "b_table".to_string(),
303 });
304 set
305 };
306
307 let value = ViewInfoValue {
308 view_info: Bytes::from([1, 2, 3].as_ref()),
309 version: 1,
310 table_names,
311 columns: vec!["a".to_string()],
312 plan_columns: vec!["number".to_string()],
313 definition: "CREATE VIEW test AS SELECT * FROM numbers".to_string(),
314 };
315 let serialized = value.try_as_raw_value().unwrap();
316 let deserialized = ViewInfoValue::try_from_raw_value(&serialized).unwrap();
317 assert_eq!(value, deserialized);
318 }
319
320 #[test]
321 fn test_update_replaces_fields_and_increments_version() {
322 let table_names = |table_name: &str| {
323 let mut set = HashSet::new();
324 set.insert(TableName {
325 catalog_name: "greptime".to_string(),
326 schema_name: "public".to_string(),
327 table_name: table_name.to_string(),
328 });
329 set
330 };
331
332 let view_info = Bytes::from([1, 2, 3].as_ref());
333 let names = table_names("a_table");
334 let columns = vec!["a".to_string()];
335 let plan_columns = vec!["number".to_string()];
336 let definition = "CREATE VIEW test AS SELECT * FROM a_table".to_string();
337 let original = ViewInfoValue::new(
338 view_info.clone(),
339 names.clone(),
340 columns.clone(),
341 plan_columns.clone(),
342 definition.clone(),
343 );
344
345 let new_view_info = Bytes::from([4, 5, 6].as_ref());
346 let new_names = table_names("b_table");
347 let new_columns = vec!["b".to_string()];
348 let new_plan_columns = vec!["number2".to_string()];
349 let new_definition = "CREATE VIEW test AS SELECT * FROM b_table".to_string();
350 let updated = original.update(
351 new_view_info.clone(),
352 new_names.clone(),
353 new_columns.clone(),
354 new_plan_columns.clone(),
355 new_definition.clone(),
356 );
357
358 assert_eq!(updated.view_info, new_view_info);
359 assert_eq!(updated.table_names, new_names);
360 assert_eq!(updated.columns, new_columns);
361 assert_eq!(updated.plan_columns, new_plan_columns);
362 assert_eq!(updated.definition, new_definition);
363 assert_eq!(original.version, 0);
364 assert_eq!(updated.version, original.version + 1);
365
366 assert_eq!(original.view_info, view_info);
367 assert_eq!(original.table_names, names);
368 assert_eq!(original.columns, columns);
369 assert_eq!(original.plan_columns, plan_columns);
370 assert_eq!(original.definition, definition);
371
372 let newest = updated.update(
373 new_view_info.clone(),
374 new_names.clone(),
375 new_columns.clone(),
376 new_plan_columns.clone(),
377 new_definition.clone(),
378 );
379 assert_eq!(newest.version, updated.version + 1);
380 assert_eq!(newest.version, 2);
381
382 let serialized = newest.try_as_raw_value().unwrap();
383 let deserialized = ViewInfoValue::try_from_raw_value(&serialized).unwrap();
384 assert_eq!(newest, deserialized);
385 }
386
387 #[test]
388 fn test_deserialize_view_info_value_with_vec_u8() {
389 #[derive(Serialize)]
390 struct OldViewInfoValue {
391 view_info: Vec<u8>,
392 table_names: HashSet<TableName>,
393 columns: Vec<String>,
394 plan_columns: Vec<String>,
395 definition: String,
396 version: u64,
397 }
398
399 let table_names = {
400 let mut set = HashSet::new();
401 set.insert(TableName {
402 catalog_name: "greptime".to_string(),
403 schema_name: "public".to_string(),
404 table_name: "a_table".to_string(),
405 });
406 set.insert(TableName {
407 catalog_name: "greptime".to_string(),
408 schema_name: "public".to_string(),
409 table_name: "b_table".to_string(),
410 });
411 set
412 };
413
414 let old_value = OldViewInfoValue {
415 view_info: vec![1, 2, 3],
416 table_names: table_names.clone(),
417 columns: vec!["a".to_string()],
418 plan_columns: vec!["number".to_string()],
419 definition: "CREATE VIEW test AS SELECT * FROM numbers".to_string(),
420 version: 1,
421 };
422
423 let serialized = serde_json::to_vec(&old_value).unwrap();
424 let deserialized = ViewInfoValue::try_from_raw_value(&serialized).unwrap();
425
426 assert_eq!(deserialized.view_info, vec![1, 2, 3]);
427 assert_eq!(deserialized.table_names, table_names);
428 assert_eq!(deserialized.columns, vec!["a".to_string()]);
429 assert_eq!(deserialized.plan_columns, vec!["number".to_string()]);
430 assert_eq!(
431 deserialized.definition,
432 "CREATE VIEW test AS SELECT * FROM numbers"
433 );
434 }
435}