common_macro/row/
schema.rs1use greptime_proto::v1::column_data_type_extension::TypeExt;
16use proc_macro2::TokenStream as TokenStream2;
17use quote::quote;
18use syn::spanned::Spanned;
19use syn::{DeriveInput, Result};
20
21use crate::row::utils::{
22 ColumnDataTypeWithExtension, ParsedField, convert_semantic_type_to_proto_semantic_type,
23 extract_struct_fields, get_column_data_type, parse_fields_from_fields_named,
24};
25use crate::row::{META_KEY_COL, META_KEY_DATATYPE};
26
27pub(crate) fn derive_schema_impl(input: DeriveInput) -> Result<TokenStream2> {
28 let Some(fields) = extract_struct_fields(&input.data) else {
29 return Err(syn::Error::new(
30 input.span(),
31 "Schema can only be derived for structs",
32 ));
33 };
34 let ident = input.ident;
35 let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
36 let fields = parse_fields_from_fields_named(fields)?;
37
38 let impl_schema_method = impl_schema_method(&fields)?;
40 Ok(quote! {
41 impl #impl_generics #ident #ty_generics #where_clause {
42 #impl_schema_method
43 }
44 })
45}
46
47fn impl_schema_method(fields: &[ParsedField<'_>]) -> Result<TokenStream2> {
48 let schemas: Vec<TokenStream2> = fields
49 .iter()
50 .map(|field| {
51 let ParsedField{ ident, column_data_type, column_attribute, ..} = field;
52 let Some(ColumnDataTypeWithExtension{data_type, extension}) = get_column_data_type(column_data_type, column_attribute)
53 else {
54 return Err(syn::Error::new(
55 ident.span(),
56 format!(
57 "expected to set data type explicitly via [({META_KEY_COL}({META_KEY_DATATYPE} = \"...\"))]"
58 ),
59 ));
60 };
61 let name = column_attribute
63 .name
64 .clone()
65 .unwrap_or_else(|| ident.to_string());
66 let name = syn::LitStr::new(&name, ident.span());
67 let column_data_type =
68 syn::LitInt::new(&(data_type as i32).to_string(), ident.span());
69 let semantic_type_val = convert_semantic_type_to_proto_semantic_type(column_attribute.semantic_type) as i32;
70 let semantic_type = syn::LitInt::new(&semantic_type_val.to_string(), ident.span());
71 let extension = match extension {
72 Some(ext) => {
73 match ext.type_ext {
74 Some(TypeExt::DecimalType(ext)) => {
75 let precision = syn::LitInt::new(&ext.precision.to_string(), ident.span());
76 let scale = syn::LitInt::new(&ext.scale.to_string(), ident.span());
77 quote! {
78 Some(ColumnDataTypeExtension { type_ext: Some(TypeExt::DecimalType(DecimalTypeExtension { precision: #precision, scale: #scale })) })
79 }
80 }
81 Some(TypeExt::JsonType(ext)) => {
82 let json_type = syn::LitInt::new(&ext.to_string(), ident.span());
83 quote! {
84 Some(ColumnDataTypeExtension { type_ext: Some(TypeExt::JsonType(#json_type)) })
85 }
86 }
87 Some(TypeExt::VectorType(ext)) => {
88 let dim = syn::LitInt::new(&ext.dim.to_string(), ident.span());
89 quote! {
90 Some(ColumnDataTypeExtension { type_ext: Some(TypeExt::VectorType(VectorTypeExtension { dim: #dim })) })
91 }
92 }
93 Some(TypeExt::ListType(ext)) => {
94 let item_type = syn::Ident::new(&ext.datatype.to_string(), ident.span());
95 quote! {
96 Some(ColumnDataTypeExtension { type_ext: Some(TypeExt::ListType(ListTypeExtension { item_type: #item_type })) })
97 }
98 }
99 Some(TypeExt::StructType(ext)) => {
100 let fields = ext.fields.iter().map(|field| {
101 let field_name = syn::Ident::new(&field.name.clone(), ident.span());
102 let field_type = syn::Ident::new(&field.datatype.to_string(), ident.span());
103 quote! {
104 StructField { name: #field_name, type_: #field_type }
105 }
106 }).collect::<Vec<_>>();
107 quote! {
108 Some(ColumnDataTypeExtension { type_ext: Some(TypeExt::StructType(StructTypeExtension { fields: [#(#fields),*] })) })
109 }
110 }
111 None => {
112 quote! { None }
113 }
114 }
115 }
116 None => quote! { None },
117 };
118
119 Ok(quote! {
120 ColumnSchema {
121 column_name: #name.to_string(),
122 datatype: #column_data_type,
123 datatype_extension: #extension,
124 options: None,
125 semantic_type: #semantic_type,
126 }
127 })
128 })
129 .collect::<Result<_>>()?;
130
131 Ok(quote! {
132 pub fn schema() -> Vec<ColumnSchema> {
133 vec![ #(#schemas),* ]
134 }
135 })
136}