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::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    // Implement `schema` method.
39    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            // Uses user explicit name or field name as column name.
62            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                        None => {
94                            quote! { None }
95                        }
96                    }
97                }
98                None => quote! { None },
99            };
100
101            Ok(quote! {
102                ColumnSchema {
103                    column_name: #name.to_string(),
104                    datatype: #column_data_type,
105                    datatype_extension: #extension,
106                    options: None,
107                    semantic_type: #semantic_type,
108                }
109            })
110        })
111        .collect::<Result<_>>()?;
112
113    Ok(quote! {
114        pub fn schema() -> Vec<ColumnSchema> {
115            vec![ #(#schemas),* ]
116        }
117    })
118}