Skip to main content

common_meta/key/
schema_name.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::{BTreeMap, HashMap};
16use std::fmt::Display;
17
18use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
19use common_time::DatabaseTimeToLive;
20use futures::stream::BoxStream;
21use humantime_serde::re::humantime;
22use serde::{Deserialize, Serialize};
23use snafu::{OptionExt, ResultExt, ensure};
24
25use crate::ensure_values;
26use crate::error::{self, Error, InvalidMetadataSnafu, ParseOptionSnafu, Result};
27use crate::key::txn_helper::TxnOpGetResponseSet;
28use crate::key::{
29    DeserializedValueWithBytes, MetadataKey, SCHEMA_NAME_KEY_PATTERN, SCHEMA_NAME_KEY_PREFIX,
30};
31use crate::kv_backend::KvBackendRef;
32use crate::kv_backend::txn::Txn;
33use crate::range_stream::{DEFAULT_PAGE_SIZE, PaginationStream};
34use crate::rpc::KeyValue;
35use crate::rpc::store::RangeRequest;
36
37const OPT_KEY_TTL: &str = "ttl";
38
39/// The schema name key, indices all schema names belong to the {catalog_name}
40///
41/// The layout:  `__schema_name/{catalog_name}/{schema_name}`.
42#[derive(Debug, Clone, Copy, PartialEq)]
43pub struct SchemaNameKey<'a> {
44    pub catalog: &'a str,
45    pub schema: &'a str,
46}
47
48impl Default for SchemaNameKey<'_> {
49    fn default() -> Self {
50        Self {
51            catalog: DEFAULT_CATALOG_NAME,
52            schema: DEFAULT_SCHEMA_NAME,
53        }
54    }
55}
56
57#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
58pub struct SchemaNameValue {
59    #[serde(default)]
60    pub ttl: Option<DatabaseTimeToLive>,
61    #[serde(default)]
62    pub extra_options: BTreeMap<String, String>,
63    /// Identifies the create-database procedure that wrote this value.
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub create_procedure_id: Option<String>,
66}
67
68impl Display for SchemaNameValue {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        if let Some(ttl) = self.ttl.map(|i| i.to_string()) {
71            writeln!(f, "'ttl'='{}'", ttl)?;
72        }
73        for (k, v) in self.extra_options.iter() {
74            writeln!(f, "'{k}'='{v}'")?;
75        }
76
77        Ok(())
78    }
79}
80
81impl TryFrom<&HashMap<String, String>> for SchemaNameValue {
82    type Error = Error;
83
84    fn try_from(value: &HashMap<String, String>) -> std::result::Result<Self, Self::Error> {
85        let ttl = value
86            .get(OPT_KEY_TTL)
87            .map(|ttl_str| {
88                ttl_str.parse::<humantime::Duration>().map_err(|_| {
89                    ParseOptionSnafu {
90                        key: OPT_KEY_TTL,
91                        value: ttl_str.clone(),
92                    }
93                    .build()
94                })
95            })
96            .transpose()?
97            .map(|ttl| ttl.into());
98        let extra_options = value
99            .iter()
100            .filter_map(|(k, v)| {
101                if k == OPT_KEY_TTL {
102                    None
103                } else {
104                    Some((k.clone(), v.clone()))
105                }
106            })
107            .collect();
108
109        Ok(Self {
110            ttl,
111            extra_options,
112            ..Default::default()
113        })
114    }
115}
116
117impl From<SchemaNameValue> for HashMap<String, String> {
118    fn from(value: SchemaNameValue) -> Self {
119        let mut opts = HashMap::new();
120        if let Some(ttl) = value.ttl.map(|ttl| ttl.to_string()) {
121            opts.insert(OPT_KEY_TTL.to_string(), ttl);
122        }
123        opts.extend(
124            value
125                .extra_options
126                .iter()
127                .map(|(k, v)| (k.clone(), v.clone())),
128        );
129        opts
130    }
131}
132
133impl<'a> SchemaNameKey<'a> {
134    pub fn new(catalog: &'a str, schema: &'a str) -> Self {
135        Self { catalog, schema }
136    }
137
138    pub fn range_start_key(catalog: &str) -> String {
139        format!("{}/{}/", SCHEMA_NAME_KEY_PREFIX, catalog)
140    }
141}
142
143impl Display for SchemaNameKey<'_> {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        write!(
146            f,
147            "{}/{}/{}",
148            SCHEMA_NAME_KEY_PREFIX, self.catalog, self.schema
149        )
150    }
151}
152
153impl<'a> MetadataKey<'a, SchemaNameKey<'a>> for SchemaNameKey<'_> {
154    fn to_bytes(&self) -> Vec<u8> {
155        self.to_string().into_bytes()
156    }
157
158    fn from_bytes(bytes: &'a [u8]) -> Result<SchemaNameKey<'a>> {
159        let key = std::str::from_utf8(bytes).map_err(|e| {
160            InvalidMetadataSnafu {
161                err_msg: format!(
162                    "SchemaNameKey '{}' is not a valid UTF8 string: {e}",
163                    String::from_utf8_lossy(bytes)
164                ),
165            }
166            .build()
167        })?;
168        SchemaNameKey::try_from(key)
169    }
170}
171
172/// Decodes `KeyValue` to {schema}
173pub fn schema_decoder(kv: KeyValue) -> Result<String> {
174    let str = std::str::from_utf8(&kv.key).context(error::ConvertRawKeySnafu)?;
175    let schema_name = SchemaNameKey::try_from(str)?;
176
177    Ok(schema_name.schema.to_string())
178}
179
180impl<'a> TryFrom<&'a str> for SchemaNameKey<'a> {
181    type Error = Error;
182
183    fn try_from(s: &'a str) -> Result<Self> {
184        let captures = SCHEMA_NAME_KEY_PATTERN
185            .captures(s)
186            .context(InvalidMetadataSnafu {
187                err_msg: format!("Illegal SchemaNameKey format: '{s}'"),
188            })?;
189
190        // Safety: pass the regex check above
191        Ok(Self {
192            catalog: captures.get(1).unwrap().as_str(),
193            schema: captures.get(2).unwrap().as_str(),
194        })
195    }
196}
197
198#[derive(Clone)]
199pub struct SchemaManager {
200    kv_backend: KvBackendRef,
201}
202
203pub type SchemaNameDecodeResult = Result<Option<DeserializedValueWithBytes<SchemaNameValue>>>;
204
205impl SchemaManager {
206    pub fn new(kv_backend: KvBackendRef) -> Self {
207        Self { kv_backend }
208    }
209
210    /// Creates `SchemaNameKey`.
211    pub async fn create(
212        &self,
213        schema: SchemaNameKey<'_>,
214        value: Option<SchemaNameValue>,
215        if_not_exists: bool,
216    ) -> Result<()> {
217        let _timer = crate::metrics::METRIC_META_CREATE_SCHEMA.start_timer();
218
219        let raw_key = schema.to_bytes();
220        let raw_value = value.unwrap_or_default().try_as_raw_value()?;
221        if self
222            .kv_backend
223            .put_conditionally(raw_key, raw_value, if_not_exists)
224            .await?
225        {
226            crate::metrics::METRIC_META_CREATE_SCHEMA_COUNTER.inc();
227        }
228
229        Ok(())
230    }
231
232    pub async fn exists(&self, schema: SchemaNameKey<'_>) -> Result<bool> {
233        let raw_key = schema.to_bytes();
234
235        self.kv_backend.exists(&raw_key).await
236    }
237
238    pub async fn get(
239        &self,
240        schema: SchemaNameKey<'_>,
241    ) -> Result<Option<DeserializedValueWithBytes<SchemaNameValue>>> {
242        let raw_key = schema.to_bytes();
243        self.kv_backend
244            .get(&raw_key)
245            .await?
246            .map(|x| DeserializedValueWithBytes::from_inner_slice(&x.value))
247            .transpose()
248    }
249
250    /// Deletes a [SchemaNameKey].
251    pub async fn delete(&self, schema: SchemaNameKey<'_>) -> Result<()> {
252        let raw_key = schema.to_bytes();
253        self.kv_backend.delete(&raw_key, false).await?;
254
255        Ok(())
256    }
257
258    pub(crate) fn build_update_txn(
259        &self,
260        schema: SchemaNameKey<'_>,
261        current_schema_value: &DeserializedValueWithBytes<SchemaNameValue>,
262        new_schema_value: &SchemaNameValue,
263    ) -> Result<(
264        Txn,
265        impl FnOnce(&mut TxnOpGetResponseSet) -> SchemaNameDecodeResult,
266    )> {
267        let raw_key = schema.to_bytes();
268        let raw_value = current_schema_value.get_raw_bytes();
269        let new_raw_value: Vec<u8> = new_schema_value.try_as_raw_value()?;
270
271        let txn = Txn::compare_and_put(raw_key.clone(), raw_value, new_raw_value);
272
273        Ok((
274            txn,
275            TxnOpGetResponseSet::decode_with(TxnOpGetResponseSet::filter(raw_key)),
276        ))
277    }
278
279    /// Updates a [SchemaNameKey].
280    pub async fn update(
281        &self,
282        schema: SchemaNameKey<'_>,
283        current_schema_value: &DeserializedValueWithBytes<SchemaNameValue>,
284        new_schema_value: &SchemaNameValue,
285    ) -> Result<()> {
286        let (txn, on_failure) =
287            self.build_update_txn(schema, current_schema_value, new_schema_value)?;
288        let mut r = self.kv_backend.txn(txn).await?;
289
290        if !r.succeeded {
291            let mut set = TxnOpGetResponseSet::from(&mut r.responses);
292            let remote_schema_value = on_failure(&mut set)?
293                .context(error::UnexpectedSnafu {
294                    err_msg:
295                        "Reads the empty schema name value in comparing operation of updating schema name value",
296                })?
297                .into_inner();
298
299            let op_name = "the updating schema name value";
300            ensure_values!(&remote_schema_value, new_schema_value, op_name);
301        }
302
303        Ok(())
304    }
305
306    /// Returns a schema stream, it lists all schemas belong to the target `catalog`.
307    pub fn schema_names(&self, catalog: &str) -> BoxStream<'static, Result<String>> {
308        let start_key = SchemaNameKey::range_start_key(catalog);
309        let req = RangeRequest::new().with_prefix(start_key.as_bytes());
310
311        let stream = PaginationStream::new(
312            self.kv_backend.clone(),
313            req,
314            DEFAULT_PAGE_SIZE,
315            schema_decoder,
316        )
317        .into_stream();
318
319        Box::pin(stream)
320    }
321}
322
323#[derive(Debug, Clone, Hash, Eq, PartialEq, Deserialize, Serialize)]
324pub struct SchemaName {
325    pub catalog_name: String,
326    pub schema_name: String,
327}
328
329impl<'a> From<&'a SchemaName> for SchemaNameKey<'a> {
330    fn from(value: &'a SchemaName) -> Self {
331        Self {
332            catalog: &value.catalog_name,
333            schema: &value.schema_name,
334        }
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use std::sync::Arc;
341    use std::time::Duration;
342
343    use super::*;
344    use crate::kv_backend::memory::MemoryKvBackend;
345
346    #[test]
347    fn test_display_schema_value() {
348        let schema_value = SchemaNameValue {
349            ttl: None,
350            ..Default::default()
351        };
352        assert_eq!("", schema_value.to_string());
353
354        let schema_value = SchemaNameValue {
355            ttl: Some(Duration::from_secs(9).into()),
356            ..Default::default()
357        };
358        assert_eq!("'ttl'='9s'\n", schema_value.to_string());
359
360        let schema_value = SchemaNameValue {
361            ttl: Some(Duration::from_secs(0).into()),
362            ..Default::default()
363        };
364        assert_eq!("'ttl'='forever'\n", schema_value.to_string());
365    }
366
367    #[test]
368    fn test_serialization() {
369        let key = SchemaNameKey::new("my-catalog", "my-schema");
370        assert_eq!(key.to_string(), "__schema_name/my-catalog/my-schema");
371
372        let parsed = SchemaNameKey::from_bytes(b"__schema_name/my-catalog/my-schema").unwrap();
373
374        assert_eq!(key, parsed);
375
376        let value = SchemaNameValue {
377            ttl: Some(Duration::from_secs(10).into()),
378            ..Default::default()
379        };
380        let mut opts: HashMap<String, String> = HashMap::new();
381        opts.insert("ttl".to_string(), "10s".to_string());
382        let from_value = SchemaNameValue::try_from(&opts).unwrap();
383        assert_eq!(value, from_value);
384
385        let parsed = SchemaNameValue::try_from_raw_value(
386            serde_json::json!({"ttl": "10s"}).to_string().as_bytes(),
387        )
388        .unwrap();
389        assert_eq!(Some(value), parsed);
390
391        let forever = SchemaNameValue {
392            ttl: Some(Default::default()),
393            ..Default::default()
394        };
395        let parsed = SchemaNameValue::try_from_raw_value(
396            serde_json::json!({"ttl": "forever"}).to_string().as_bytes(),
397        )
398        .unwrap();
399        assert_eq!(Some(forever), parsed);
400
401        let instant_err = SchemaNameValue::try_from_raw_value(
402            serde_json::json!({"ttl": "instant"}).to_string().as_bytes(),
403        );
404        assert!(instant_err.is_err());
405
406        let none = SchemaNameValue::try_from_raw_value("null".as_bytes()).unwrap();
407        assert!(none.is_none());
408
409        let err_empty = SchemaNameValue::try_from_raw_value("".as_bytes());
410        assert!(err_empty.is_err());
411    }
412
413    #[test]
414    fn test_extra_options_compatibility() {
415        // Test with extra_options only
416        let mut opts: HashMap<String, String> = HashMap::new();
417        opts.insert("foo".to_string(), "bar".to_string());
418        opts.insert("baz".to_string(), "qux".to_string());
419        let value = SchemaNameValue::try_from(&opts).unwrap();
420        assert_eq!(value.ttl, None);
421        assert_eq!(value.extra_options.get("foo"), Some(&"bar".to_string()));
422        assert_eq!(value.extra_options.get("baz"), Some(&"qux".to_string()));
423
424        // Test round-trip conversion
425        let opts_back: HashMap<String, String> = value.clone().into();
426        assert_eq!(opts_back.get("foo"), Some(&"bar".to_string()));
427        assert_eq!(opts_back.get("baz"), Some(&"qux".to_string()));
428        assert!(!opts_back.contains_key("ttl"));
429
430        // Test with both ttl and extra_options
431        let mut opts: HashMap<String, String> = HashMap::new();
432        opts.insert("ttl".to_string(), "5m".to_string());
433        opts.insert("opt1".to_string(), "val1".to_string());
434        let value = SchemaNameValue::try_from(&opts).unwrap();
435        assert_eq!(value.ttl, Some(Duration::from_secs(300).into()));
436        assert_eq!(value.extra_options.get("opt1"), Some(&"val1".to_string()));
437
438        // Test serialization/deserialization compatibility
439        let json = serde_json::to_string(&value).unwrap();
440        let deserialized: SchemaNameValue = serde_json::from_str(&json).unwrap();
441        assert_eq!(value, deserialized);
442
443        // Test display includes extra_options
444        let mut value = SchemaNameValue::default();
445        value
446            .extra_options
447            .insert("foo".to_string(), "bar".to_string());
448        let display = value.to_string();
449        assert!(display.contains("'foo'='bar'"));
450    }
451
452    #[test]
453    fn test_backward_compatibility_with_old_format() {
454        // Simulate old format: only ttl, no extra_options
455        let json = r#"{"ttl":"10s"}"#;
456        let parsed = SchemaNameValue::try_from_raw_value(json.as_bytes()).unwrap();
457        assert_eq!(
458            parsed,
459            Some(SchemaNameValue {
460                ttl: Some(Duration::from_secs(10).into()),
461                extra_options: BTreeMap::new(),
462                create_procedure_id: None,
463            })
464        );
465
466        // Simulate old format: null value
467        let json = r#"null"#;
468        let parsed = SchemaNameValue::try_from_raw_value(json.as_bytes()).unwrap();
469        assert!(parsed.is_none());
470    }
471
472    #[test]
473    fn test_forward_compatibility_with_new_options() {
474        // Simulate new format: ttl + extra_options
475        let json = r#"{"ttl":"15s","extra_options":{"foo":"bar","baz":"qux"}}"#;
476        let parsed = SchemaNameValue::try_from_raw_value(json.as_bytes()).unwrap();
477        let mut expected_options = BTreeMap::new();
478        expected_options.insert("foo".to_string(), "bar".to_string());
479        expected_options.insert("baz".to_string(), "qux".to_string());
480        assert_eq!(
481            parsed,
482            Some(SchemaNameValue {
483                ttl: Some(Duration::from_secs(15).into()),
484                extra_options: expected_options,
485                create_procedure_id: None,
486            })
487        );
488    }
489
490    #[test]
491    fn test_create_procedure_id_serialization() {
492        let value = SchemaNameValue {
493            create_procedure_id: Some("4ee0ba94-11f0-4d4d-9468-5ebf732e3ab2".to_string()),
494            ..Default::default()
495        };
496        let raw = value.try_as_raw_value().unwrap();
497        assert_eq!(
498            SchemaNameValue::try_from_raw_value(&raw).unwrap(),
499            Some(value)
500        );
501
502        let raw = SchemaNameValue::default().try_as_raw_value().unwrap();
503        assert!(
504            !String::from_utf8(raw)
505                .unwrap()
506                .contains("create_procedure_id")
507        );
508    }
509
510    #[tokio::test]
511    async fn test_key_exist() {
512        let manager = SchemaManager::new(Arc::new(MemoryKvBackend::default()));
513        let schema_key = SchemaNameKey::new("my-catalog", "my-schema");
514        manager.create(schema_key, None, false).await.unwrap();
515
516        assert!(manager.exists(schema_key).await.unwrap());
517
518        let wrong_schema_key = SchemaNameKey::new("my-catalog", "my-wrong");
519
520        assert!(!manager.exists(wrong_schema_key).await.unwrap());
521    }
522
523    #[tokio::test]
524    async fn test_update_schema_value() {
525        let manager = SchemaManager::new(Arc::new(MemoryKvBackend::default()));
526        let schema_key = SchemaNameKey::new("my-catalog", "my-schema");
527        manager.create(schema_key, None, false).await.unwrap();
528
529        let current_schema_value = manager.get(schema_key).await.unwrap().unwrap();
530        let new_schema_value = SchemaNameValue {
531            ttl: Some(Duration::from_secs(10).into()),
532            ..Default::default()
533        };
534        manager
535            .update(schema_key, &current_schema_value, &new_schema_value)
536            .await
537            .unwrap();
538
539        // Update with the same value, should be ok
540        manager
541            .update(schema_key, &current_schema_value, &new_schema_value)
542            .await
543            .unwrap();
544
545        let new_schema_value = SchemaNameValue {
546            ttl: Some(Duration::from_secs(40).into()),
547            ..Default::default()
548        };
549        let incorrect_schema_value = SchemaNameValue {
550            ttl: Some(Duration::from_secs(20).into()),
551            ..Default::default()
552        }
553        .try_as_raw_value()
554        .unwrap();
555        let incorrect_schema_value =
556            DeserializedValueWithBytes::from_inner_slice(&incorrect_schema_value).unwrap();
557
558        manager
559            .update(schema_key, &incorrect_schema_value, &new_schema_value)
560            .await
561            .unwrap_err();
562
563        let current_schema_value = manager.get(schema_key).await.unwrap().unwrap();
564        let new_schema_value = SchemaNameValue {
565            ttl: None,
566            ..Default::default()
567        };
568        manager
569            .update(schema_key, &current_schema_value, &new_schema_value)
570            .await
571            .unwrap();
572
573        let current_schema_value = manager.get(schema_key).await.unwrap().unwrap();
574        assert_eq!(new_schema_value, *current_schema_value);
575    }
576}