Skip to main content

common_meta/key/
view_info.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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/// The key stores the metadata of the view.
36///
37/// The layout: `__view_info/{view_id}`.
38#[derive(Debug, PartialEq)]
39pub struct ViewInfoKey {
40    view_id: TableId,
41}
42
43impl ViewInfoKey {
44    /// Returns a new `[ViewInfoKey]`.
45    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        // Safety: pass the regex check above
78        let view_id = captures[1].parse::<TableId>().unwrap();
79        Ok(ViewInfoKey { view_id })
80    }
81}
82
83/// The VIEW info value that keeps the metadata.
84#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
85pub struct ViewInfoValue {
86    // The encoded logical plan
87    pub view_info: Bytes,
88    // The resolved fully table names in logical plan
89    pub table_names: HashSet<TableName>,
90    // The view columns
91    pub columns: Vec<String>,
92    // The original plan columns
93    pub plan_columns: Vec<String>,
94    // The SQL to create the view
95    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    /// Builds a replacement [`ViewInfoValue`] with the next version, leaving this
118    /// value unchanged for a caller's compare-and-swap.
119    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
138/// The `[ViewInfo]` manager
139pub 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    /// Builds a create view info transaction, it expected the `__view_info/{view_id}` wasn't occupied.
153    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    /// Builds a update view info transaction, it expected the remote value equals the `current_current_view_info_value`.
173    /// It retrieves the latest value if the comparing failed.
174    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    /// Get the `[ViewInfoValue]` by the view id
197    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    /// Get the `[ViewInfoValue]` by the view id slice in batch
211    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                    // Safety: must exist.
230                    **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    /// Returns batch of `DeserializedValueWithBytes<ViewInfoValue>`.
240    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                    // Safety: must exist.
262                    **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}