1use 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#[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 #[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
172pub 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 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 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 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 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 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 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 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 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 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 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 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 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 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, ¤t_schema_value, &new_schema_value)
536 .await
537 .unwrap();
538
539 manager
541 .update(schema_key, ¤t_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, ¤t_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}