common_macro/row/
schema.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 greptime_proto::v1::ColumnDataTypeExtension;
16use greptime_proto::v1::column_data_type_extension::TypeExt;
17use proc_macro2::{Span, TokenStream as TokenStream2};
18use quote::quote;
19use syn::spanned::Spanned;
20use syn::{DeriveInput, Result};
21
22use crate::row::utils::{
23    ColumnDataTypeWithExtension, ParsedField, convert_semantic_type_to_proto_semantic_type,
24    extract_struct_fields, get_column_data_type, parse_fields_from_fields_named,
25};
26use crate::row::{META_KEY_COL, META_KEY_DATATYPE};
27
28pub(crate) fn derive_schema_impl(input: DeriveInput) -> Result<TokenStream2> {
29    let Some(fields) = extract_struct_fields(&input.data) else {
30        return Err(syn::Error::new(
31            input.span(),
32            "Schema can only be derived for structs",
33        ));
34    };
35    let ident = input.ident;
36    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
37    let fields = parse_fields_from_fields_named(fields)?;
38
39    // Implement `schema` method.
40    let impl_schema_method = impl_schema_method(&fields)?;
41    Ok(quote! {
42        impl #impl_generics #ident #ty_generics #where_clause {
43            #impl_schema_method
44        }
45    })
46}
47
48fn impl_schema_method(fields: &[ParsedField<'_>]) -> Result<TokenStream2> {
49    let schemas: Vec<TokenStream2> = fields
50        .iter()
51        .map(|field| {
52            let ParsedField{ ident, column_data_type, column_attribute, ..} = field;
53            let Some(ColumnDataTypeWithExtension{data_type, extension}) = get_column_data_type(column_data_type, column_attribute)
54            else {
55                return Err(syn::Error::new(
56                    ident.span(),
57                    format!(
58                        "expected to set data type explicitly via [({META_KEY_COL}({META_KEY_DATATYPE} = \"...\"))]"
59                    ),
60                ));
61            };
62            // Uses user explicit name or field name as column name.
63            let name = column_attribute
64                .name
65                .clone()
66                .unwrap_or_else(|| ident.to_string());
67            let name = syn::LitStr::new(&name, ident.span());
68            let column_data_type =
69                syn::LitInt::new(&(data_type as i32).to_string(), ident.span());
70            let semantic_type_val = convert_semantic_type_to_proto_semantic_type(column_attribute.semantic_type) as i32;
71            let semantic_type = syn::LitInt::new(&semantic_type_val.to_string(), ident.span());
72            let extension = match extension {
73                Some(ext) => column_data_type_extension_to_tokens(&ext, ident.span()),
74                None => quote! { None },
75            };
76
77            Ok(quote! {
78                ColumnSchema {
79                    column_name: #name.to_string(),
80                    datatype: #column_data_type,
81                    datatype_extension: #extension,
82                    options: None,
83                    semantic_type: #semantic_type,
84                }
85            })
86        })
87        .collect::<Result<_>>()?;
88
89    Ok(quote! {
90        pub fn schema() -> Vec<ColumnSchema> {
91            vec![ #(#schemas),* ]
92        }
93    })
94}
95
96fn column_data_type_extension_to_tokens(
97    extension: &ColumnDataTypeExtension,
98    span: Span,
99) -> TokenStream2 {
100    match extension.type_ext.as_ref() {
101        Some(TypeExt::DecimalType(ext)) => {
102            let precision = syn::LitInt::new(&ext.precision.to_string(), span);
103            let scale = syn::LitInt::new(&ext.scale.to_string(), span);
104            quote! {
105                Some(ColumnDataTypeExtension {
106                    type_ext: Some(TypeExt::DecimalType(DecimalTypeExtension {
107                        precision: #precision,
108                        scale: #scale,
109                    })),
110                })
111            }
112        }
113        Some(TypeExt::JsonType(ext)) => {
114            let json_type = syn::LitInt::new(&ext.to_string(), span);
115            quote! {
116                Some(ColumnDataTypeExtension {
117                    type_ext: Some(TypeExt::JsonType(#json_type)),
118                })
119            }
120        }
121        Some(TypeExt::VectorType(ext)) => {
122            let dim = syn::LitInt::new(&ext.dim.to_string(), span);
123            quote! {
124                Some(ColumnDataTypeExtension {
125                    type_ext: Some(TypeExt::VectorType(VectorTypeExtension { dim: #dim })),
126                })
127            }
128        }
129        Some(TypeExt::ListType(ext)) => {
130            let datatype = syn::LitInt::new(&ext.datatype.to_string(), span);
131            let datatype_extension = ext
132                .datatype_extension
133                .as_deref()
134                .map(|ext| column_data_type_extension_to_tokens(ext, span))
135                .unwrap_or_else(|| quote! { None });
136            quote! {
137                Some(ColumnDataTypeExtension {
138                    type_ext: Some(TypeExt::ListType(Box::new(ListTypeExtension {
139                        datatype: #datatype,
140                        datatype_extension: #datatype_extension,
141                    }))),
142                })
143            }
144        }
145        Some(TypeExt::StructType(ext)) => {
146            let fields = ext.fields.iter().map(|field| {
147                let field_name = &field.name;
148                let datatype = syn::LitInt::new(&field.datatype.to_string(), span);
149                let datatype_extension = field
150                    .datatype_extension
151                    .as_ref()
152                    .map(|ext| column_data_type_extension_to_tokens(ext, span))
153                    .unwrap_or_else(|| quote! { None });
154                quote! {
155                    greptime_proto::v1::StructField {
156                        name: #field_name.to_string(),
157                        datatype: #datatype,
158                        datatype_extension: #datatype_extension,
159                    }
160                }
161            });
162            quote! {
163                Some(ColumnDataTypeExtension {
164                    type_ext: Some(TypeExt::StructType(StructTypeExtension {
165                        fields: vec![#(#fields),*],
166                    })),
167                })
168            }
169        }
170        Some(TypeExt::JsonNativeType(ext)) => {
171            let inner = syn::LitInt::new(&ext.datatype.to_string(), span);
172            let datatype_extension = ext
173                .datatype_extension
174                .as_deref()
175                .map(|ext| column_data_type_extension_to_tokens(ext, span))
176                .unwrap_or_else(|| quote! { None });
177            quote! {
178                Some(ColumnDataTypeExtension {
179                    type_ext: Some(TypeExt::JsonNativeType(Box::new(
180                        JsonNativeTypeExtension {
181                            datatype: #inner,
182                            datatype_extension: #datatype_extension,
183                        },
184                    ))),
185                })
186            }
187        }
188        Some(TypeExt::DictionaryType(ext)) => {
189            let key_datatype = syn::LitInt::new(&ext.key_datatype.to_string(), span);
190            let value_datatype = syn::LitInt::new(&ext.value_datatype.to_string(), span);
191            let key_datatype_extension = ext
192                .key_datatype_extension
193                .as_deref()
194                .map(|ext| column_data_type_extension_to_tokens(ext, span))
195                .unwrap_or_else(|| quote! { None });
196            let value_datatype_extension = ext
197                .value_datatype_extension
198                .as_deref()
199                .map(|ext| column_data_type_extension_to_tokens(ext, span))
200                .unwrap_or_else(|| quote! { None });
201            quote! {
202                Some(ColumnDataTypeExtension {
203                    type_ext: Some(TypeExt::DictionaryType(Box::new(
204                        DictionaryTypeExtension {
205                            key_datatype: #key_datatype,
206                            key_datatype_extension: #key_datatype_extension,
207                            value_datatype: #value_datatype,
208                            value_datatype_extension: #value_datatype_extension,
209                        },
210                    ))),
211                })
212            }
213        }
214        None => quote! { None },
215    }
216}
common_macro/row/schema.rs

common_macro/row/
schema.rs
