1use std::collections::{BTreeMap, BTreeSet};
6
7use thiserror::Error;
8use ytsaurus_yson::{YsonNode, YsonValue};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub enum WireType {
13 Nothing,
15 Boolean,
17 Int8,
19 Int16,
21 Int32,
23 Int64,
25 Int128,
27 Int256,
29 Uint8,
31 Uint16,
33 Uint32,
35 Uint64,
37 Double,
39 String32,
41 Yson32,
43 Variant8,
45 Variant16,
47 RepeatedVariant8,
49 RepeatedVariant16,
51 Tuple,
53}
54
55impl WireType {
56 #[must_use]
58 pub const fn as_str(self) -> &'static str {
59 match self {
60 Self::Nothing => "nothing",
61 Self::Boolean => "boolean",
62 Self::Int8 => "int8",
63 Self::Int16 => "int16",
64 Self::Int32 => "int32",
65 Self::Int64 => "int64",
66 Self::Int128 => "int128",
67 Self::Int256 => "int256",
68 Self::Uint8 => "uint8",
69 Self::Uint16 => "uint16",
70 Self::Uint32 => "uint32",
71 Self::Uint64 => "uint64",
72 Self::Double => "double",
73 Self::String32 => "string32",
74 Self::Yson32 => "yson32",
75 Self::Variant8 => "variant8",
76 Self::Variant16 => "variant16",
77 Self::RepeatedVariant8 => "repeated_variant8",
78 Self::RepeatedVariant16 => "repeated_variant16",
79 Self::Tuple => "tuple",
80 }
81 }
82
83 #[must_use]
85 pub fn parse(value: &str) -> Option<Self> {
86 Some(match value {
87 "nothing" => Self::Nothing,
88 "boolean" => Self::Boolean,
89 "int8" => Self::Int8,
90 "int16" => Self::Int16,
91 "int32" => Self::Int32,
92 "int64" => Self::Int64,
93 "int128" => Self::Int128,
94 "int256" => Self::Int256,
95 "uint8" => Self::Uint8,
96 "uint16" => Self::Uint16,
97 "uint32" => Self::Uint32,
98 "uint64" => Self::Uint64,
99 "double" => Self::Double,
100 "string32" => Self::String32,
101 "yson32" => Self::Yson32,
102 "variant8" => Self::Variant8,
103 "variant16" => Self::Variant16,
104 "repeated_variant8" => Self::RepeatedVariant8,
105 "repeated_variant16" => Self::RepeatedVariant16,
106 "tuple" => Self::Tuple,
107 _ => return None,
108 })
109 }
110
111 #[must_use]
113 pub const fn is_simple(self) -> bool {
114 matches!(
115 self,
116 Self::Boolean
117 | Self::Int8
118 | Self::Int16
119 | Self::Int32
120 | Self::Int64
121 | Self::Int128
122 | Self::Int256
123 | Self::Uint8
124 | Self::Uint16
125 | Self::Uint32
126 | Self::Uint64
127 | Self::Double
128 | Self::String32
129 | Self::Yson32
130 )
131 }
132
133 #[must_use]
135 pub const fn fixed_width(self) -> Option<usize> {
136 match self {
137 Self::Nothing => Some(0),
138 Self::Boolean | Self::Int8 | Self::Uint8 => Some(1),
139 Self::Int16 | Self::Uint16 => Some(2),
140 Self::Int32 | Self::Uint32 => Some(4),
141 Self::Int64 | Self::Uint64 | Self::Double => Some(8),
142 Self::Int128 => Some(16),
143 Self::Int256 => Some(32),
144 Self::String32
145 | Self::Yson32
146 | Self::Variant8
147 | Self::Variant16
148 | Self::RepeatedVariant8
149 | Self::RepeatedVariant16
150 | Self::Tuple => None,
151 }
152 }
153}
154
155impl std::fmt::Display for WireType {
156 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157 formatter.write_str(self.as_str())
158 }
159}
160
161#[derive(Debug, Clone, PartialEq, Eq)]
163pub struct Schema {
164 pub wire_type: WireType,
166 pub name: Option<String>,
168 pub children: Vec<Schema>,
170}
171
172impl Schema {
173 #[must_use]
175 pub const fn leaf(wire_type: WireType) -> Self {
176 Self {
177 wire_type,
178 name: None,
179 children: Vec::new(),
180 }
181 }
182
183 #[must_use]
185 pub fn named(name: impl Into<String>, wire_type: WireType) -> Self {
186 Self {
187 wire_type,
188 name: Some(name.into()),
189 children: Vec::new(),
190 }
191 }
192
193 #[must_use]
195 pub fn tuple(children: impl IntoIterator<Item = Schema>) -> Self {
196 Self {
197 wire_type: WireType::Tuple,
198 name: None,
199 children: children.into_iter().collect(),
200 }
201 }
202
203 #[must_use]
206 pub fn optional(self) -> Self {
207 Self {
208 wire_type: WireType::Variant8,
209 name: self.name.clone(),
210 children: vec![Schema::leaf(WireType::Nothing), Self { name: None, ..self }],
211 }
212 }
213
214 pub fn validate(&self) -> Result<(), SchemaError> {
220 let count = self.children.len();
221 if self.wire_type == WireType::Nothing || self.wire_type.is_simple() {
222 if count != 0 {
223 return Err(SchemaError::UnexpectedChildren {
224 wire_type: self.wire_type,
225 count,
226 });
227 }
228 } else {
229 match self.wire_type {
230 WireType::Variant8 if count > 256 => {
231 return Err(SchemaError::TooManyChildren {
232 wire_type: self.wire_type,
233 count,
234 maximum: 256,
235 });
236 }
237 WireType::Variant16 if count > 65_536 => {
238 return Err(SchemaError::TooManyChildren {
239 wire_type: self.wire_type,
240 count,
241 maximum: 65_536,
242 });
243 }
244 WireType::RepeatedVariant8 if count > 255 => {
245 return Err(SchemaError::TooManyChildren {
246 wire_type: self.wire_type,
247 count,
248 maximum: 255,
249 });
250 }
251 WireType::RepeatedVariant16 if count > 65_535 => {
252 return Err(SchemaError::TooManyChildren {
253 wire_type: self.wire_type,
254 count,
255 maximum: 65_535,
256 });
257 }
258 WireType::Variant8
259 | WireType::Variant16
260 | WireType::RepeatedVariant8
261 | WireType::RepeatedVariant16
262 | WireType::Tuple => {}
263 _ => unreachable!("Skiff leaves were handled before compound validation"),
265 }
266 }
267
268 for child in &self.children {
269 child.validate()?;
270 }
271 Ok(())
272 }
273
274 #[must_use]
276 pub fn to_yson(&self) -> YsonValue {
277 let mut fields = BTreeMap::new();
278 fields.insert(b"wire_type".to_vec(), string(self.wire_type.as_str()));
279 if let Some(name) = &self.name {
280 fields.insert(b"name".to_vec(), string(name));
281 }
282 if !self.children.is_empty() {
283 fields.insert(
284 b"children".to_vec(),
285 list(self.children.iter().map(Self::to_yson)),
286 );
287 }
288 value(YsonNode::Map(fields))
289 }
290
291 pub fn from_yson(input: &YsonValue) -> Result<Self, SchemaError> {
293 reject_attributes(input, "schema")?;
294 let fields = map_fields(input, "schema")?;
295 reject_unknown(fields, &[b"wire_type", b"name", b"children"], "schema")?;
296
297 let wire_type = required_string(fields, b"wire_type", "schema")?;
298 let wire_type = WireType::parse(wire_type)
299 .ok_or_else(|| SchemaError::UnknownWireType(wire_type.to_owned()))?;
300 let name = optional_string(fields, b"name", "schema")?.map(str::to_owned);
301 let children = match fields.get(b"children".as_slice()) {
302 None => Vec::new(),
303 Some(child_values) => list_items(child_values, "schema.children")?
304 .iter()
305 .map(Self::from_yson)
306 .collect::<Result<_, _>>()?,
307 };
308
309 let schema = Self {
310 wire_type,
311 name,
312 children,
313 };
314 schema.validate()?;
315 Ok(schema)
316 }
317}
318
319#[derive(Debug, Clone, PartialEq, Eq)]
321pub enum SchemaRef {
322 Inline(Schema),
324 Registry(String),
326}
327
328#[derive(Debug, Clone, PartialEq, Eq)]
330pub struct Format {
331 table_schemas: Vec<SchemaRef>,
332 schema_registry: BTreeMap<String, Schema>,
333}
334
335impl Format {
336 pub fn new(table_schemas: Vec<SchemaRef>) -> Result<Self, SchemaError> {
340 Self::from_parts(table_schemas, BTreeMap::new())
341 }
342
343 pub fn from_parts(
345 table_schemas: Vec<SchemaRef>,
346 schema_registry: BTreeMap<String, Schema>,
347 ) -> Result<Self, SchemaError> {
348 let format = Self {
349 table_schemas,
350 schema_registry,
351 };
352 format.validate()?;
353 Ok(format)
354 }
355
356 #[must_use]
358 pub fn table_schemas(&self) -> &[SchemaRef] {
359 &self.table_schemas
360 }
361
362 #[must_use]
364 pub fn schema_registry(&self) -> &BTreeMap<String, Schema> {
365 &self.schema_registry
366 }
367
368 pub fn table_schema(&self, index: usize) -> Result<&Schema, SchemaError> {
370 let reference = self
371 .table_schemas
372 .get(index)
373 .ok_or(SchemaError::MissingTableSchema { index })?;
374 match reference {
375 SchemaRef::Inline(schema) => Ok(schema),
376 SchemaRef::Registry(name) => self
377 .schema_registry
378 .get(name)
379 .ok_or_else(|| SchemaError::UnknownRegistryReference(name.clone())),
380 }
381 }
382
383 pub fn validate(&self) -> Result<(), SchemaError> {
385 if self.table_schemas.is_empty() {
386 return Err(SchemaError::EmptyTableSchemas);
387 }
388 for schema in self.schema_registry.values() {
389 schema.validate()?;
390 }
391 for (index, reference) in self.table_schemas.iter().enumerate() {
392 match reference {
393 SchemaRef::Inline(schema) => schema.validate()?,
394 SchemaRef::Registry(name) if self.schema_registry.contains_key(name) => {}
395 SchemaRef::Registry(name) => {
396 return Err(SchemaError::UnknownRegistryReference(name.clone()));
397 }
398 }
399 validate_table_schema(self.table_schema(index)?)?;
402 }
403 Ok(())
404 }
405
406 #[must_use]
408 pub fn to_yson(&self) -> YsonValue {
409 let mut attributes = BTreeMap::new();
410 attributes.insert(
411 b"table_skiff_schemas".to_vec(),
412 list(self.table_schemas.iter().map(SchemaRef::to_yson)),
413 );
414 if !self.schema_registry.is_empty() {
415 let mut registry = BTreeMap::new();
416 for (name, schema) in &self.schema_registry {
417 registry.insert(name.as_bytes().to_vec(), schema.to_yson());
418 }
419 attributes.insert(
420 b"skiff_schema_registry".to_vec(),
421 value(YsonNode::Map(registry)),
422 );
423 }
424 YsonValue {
425 attributes: Some(attributes),
426 node: YsonNode::String(b"skiff".to_vec()),
427 }
428 }
429
430 pub fn from_yson(input: &YsonValue) -> Result<Self, SchemaError> {
432 let YsonNode::String(name) = &input.node else {
433 return Err(SchemaError::FormatMustBeSkiff);
434 };
435 if name.as_slice() != b"skiff" {
436 return Err(SchemaError::FormatMustBeSkiff);
437 }
438 let attributes = input
439 .attributes
440 .as_ref()
441 .ok_or(SchemaError::MissingFormatAttribute("table_skiff_schemas"))?;
442 reject_unknown(
443 attributes,
444 &[b"table_skiff_schemas", b"skiff_schema_registry"],
445 "format attributes",
446 )?;
447
448 let table_values = attributes
449 .get(b"table_skiff_schemas".as_slice())
450 .ok_or(SchemaError::MissingFormatAttribute("table_skiff_schemas"))?;
451 let table_schemas = list_items(table_values, "format.table_skiff_schemas")?
452 .iter()
453 .map(SchemaRef::from_yson)
454 .collect::<Result<_, _>>()?;
455
456 let schema_registry = match attributes.get(b"skiff_schema_registry".as_slice()) {
457 None => BTreeMap::new(),
458 Some(value) => {
459 reject_attributes(value, "format.skiff_schema_registry")?;
460 let entries = map_fields(value, "format.skiff_schema_registry")?;
461 let mut registry = BTreeMap::new();
462 for (name, schema) in entries {
463 let name = std::str::from_utf8(name).map_err(|_| SchemaError::InvalidUtf8 {
464 field: "format.skiff_schema_registry key",
465 })?;
466 registry.insert(name.to_owned(), Schema::from_yson(schema)?);
467 }
468 registry
469 }
470 };
471
472 Self::from_parts(table_schemas, schema_registry)
473 }
474}
475
476impl SchemaRef {
477 fn to_yson(&self) -> YsonValue {
478 match self {
479 Self::Inline(schema) => schema.to_yson(),
480 Self::Registry(name) => string(format!("${name}")),
481 }
482 }
483
484 fn from_yson(input: &YsonValue) -> Result<Self, SchemaError> {
485 reject_attributes(input, "table schema reference")?;
486 match &input.node {
487 YsonNode::String(name) if name.first() == Some(&b'$') && name.len() > 1 => {
488 let name =
489 std::str::from_utf8(&name[1..]).map_err(|_| SchemaError::InvalidUtf8 {
490 field: "table schema registry reference",
491 })?;
492 Ok(Self::Registry(name.to_owned()))
493 }
494 YsonNode::String(_) => Err(SchemaError::InvalidRegistryReference),
495 YsonNode::Map(_) => Ok(Self::Inline(Schema::from_yson(input)?)),
496 _ => Err(SchemaError::InvalidSchemaReference),
497 }
498 }
499}
500
501#[derive(Debug, Error, Clone, PartialEq, Eq)]
503pub enum SchemaError {
504 #[error("unknown Skiff wire type {0:?}")]
506 UnknownWireType(String),
507 #[error("Skiff {wire_type} cannot have {count} child schema node(s)")]
509 UnexpectedChildren {
510 wire_type: WireType,
512 count: usize,
514 },
515 #[error("Skiff {wire_type} has {count} children, exceeding its {maximum} child limit")]
517 TooManyChildren {
518 wire_type: WireType,
520 count: usize,
522 maximum: usize,
524 },
525 #[error("Skiff format requires at least one table schema")]
527 EmptyTableSchemas,
528 #[error("Skiff schema registry has no entry named {0:?}")]
530 UnknownRegistryReference(String),
531 #[error("Skiff format has no schema for table index {index}")]
533 MissingTableSchema {
534 index: usize,
536 },
537 #[error("Skiff format must be the attributed string \"skiff\"")]
539 FormatMustBeSkiff,
540 #[error("Skiff format is missing required attribute {0:?}")]
542 MissingFormatAttribute(&'static str),
543 #[error("Skiff registry reference must be a non-empty string beginning with '$'")]
545 InvalidRegistryReference,
546 #[error("Skiff table schema reference must be a schema map or '$' registry reference")]
548 InvalidSchemaReference,
549 #[error("Skiff {context} must not carry YSON attributes")]
551 UnexpectedAttributes {
552 context: &'static str,
554 },
555 #[error("Skiff {context} must be a YSON map")]
557 ExpectedMap {
558 context: &'static str,
560 },
561 #[error("Skiff {context} must be a YSON list")]
563 ExpectedList {
564 context: &'static str,
566 },
567 #[error("Skiff {context} is missing required field {field:?}")]
569 MissingField {
570 context: &'static str,
572 field: &'static str,
574 },
575 #[error("Skiff {context}.{field} must be a YSON string")]
577 ExpectedString {
578 context: &'static str,
580 field: &'static str,
582 },
583 #[error("Skiff {field} must be valid UTF-8")]
585 InvalidUtf8 {
586 field: &'static str,
588 },
589 #[error("Skiff {context} has unsupported field {field:?}")]
591 UnknownField {
592 context: &'static str,
594 field: String,
596 },
597 #[error("Skiff table schema root must be tuple, got {found}")]
599 TableSchemaRootMustBeTuple {
600 found: WireType,
602 },
603 #[error("Skiff table schema child {index} must have a non-empty name")]
605 TableSchemaChildMissingName {
606 index: usize,
608 },
609}
610
611pub(crate) fn validate_table_schema(schema: &Schema) -> Result<(), SchemaError> {
612 if schema.wire_type != WireType::Tuple {
613 return Err(SchemaError::TableSchemaRootMustBeTuple {
614 found: schema.wire_type,
615 });
616 }
617 for (index, child) in schema.children.iter().enumerate() {
618 if child.name.as_deref().is_none_or(str::is_empty) {
619 return Err(SchemaError::TableSchemaChildMissingName { index });
620 }
621 }
622 Ok(())
623}
624
625fn string(bytes: impl AsRef<[u8]>) -> YsonValue {
626 value(YsonNode::String(bytes.as_ref().to_vec()))
627}
628
629fn list(values: impl IntoIterator<Item = YsonValue>) -> YsonValue {
630 value(YsonNode::List(values.into_iter().collect()))
631}
632
633fn value(node: YsonNode) -> YsonValue {
634 YsonValue {
635 attributes: None,
636 node,
637 }
638}
639
640fn reject_attributes(value: &YsonValue, context: &'static str) -> Result<(), SchemaError> {
641 if value.attributes.is_some() {
642 return Err(SchemaError::UnexpectedAttributes { context });
643 }
644 Ok(())
645}
646
647fn map_fields<'a>(
648 value: &'a YsonValue,
649 context: &'static str,
650) -> Result<&'a BTreeMap<Vec<u8>, YsonValue>, SchemaError> {
651 match &value.node {
652 YsonNode::Map(fields) => Ok(fields),
653 _ => Err(SchemaError::ExpectedMap { context }),
654 }
655}
656
657fn list_items<'a>(
658 value: &'a YsonValue,
659 context: &'static str,
660) -> Result<&'a [YsonValue], SchemaError> {
661 match &value.node {
662 YsonNode::List(items) => Ok(items),
663 _ => Err(SchemaError::ExpectedList { context }),
664 }
665}
666
667fn required_string<'a>(
668 fields: &'a BTreeMap<Vec<u8>, YsonValue>,
669 field: &'static [u8],
670 context: &'static str,
671) -> Result<&'a str, SchemaError> {
672 let value = fields.get(field).ok_or(SchemaError::MissingField {
673 context,
674 field: std::str::from_utf8(field).expect("literal field name"),
675 })?;
676 as_utf8_string(
677 value,
678 context,
679 std::str::from_utf8(field).expect("literal field name"),
680 )
681}
682
683fn optional_string<'a>(
684 fields: &'a BTreeMap<Vec<u8>, YsonValue>,
685 field: &'static [u8],
686 context: &'static str,
687) -> Result<Option<&'a str>, SchemaError> {
688 let field_name = std::str::from_utf8(field).expect("literal field name");
689 fields
690 .get(field)
691 .map(|value| as_utf8_string(value, context, field_name))
692 .transpose()
693}
694
695fn as_utf8_string<'a>(
696 value: &'a YsonValue,
697 context: &'static str,
698 field: &'static str,
699) -> Result<&'a str, SchemaError> {
700 reject_attributes(value, context)?;
701 let YsonNode::String(bytes) = &value.node else {
702 return Err(SchemaError::ExpectedString { context, field });
703 };
704 std::str::from_utf8(bytes).map_err(|_| SchemaError::InvalidUtf8 { field })
705}
706
707fn reject_unknown(
708 fields: &BTreeMap<Vec<u8>, YsonValue>,
709 allowed: &[&[u8]],
710 context: &'static str,
711) -> Result<(), SchemaError> {
712 let allowed = allowed.iter().copied().collect::<BTreeSet<_>>();
713 for field in fields.keys() {
714 if !allowed.contains(field.as_slice()) {
715 return Err(SchemaError::UnknownField {
716 context,
717 field: String::from_utf8_lossy(field).into_owned(),
718 });
719 }
720 }
721 Ok(())
722}