use proc_macro::TokenStream;
use quote::{quote, format_ident};
use syn::{
parse_macro_input, Data, DeriveInput, Fields, Attribute, Lit,
};
struct TableAttrs {
name: String,
_engine: Option<String>,
_partition_by: Option<String>,
_order_by: Option<String>,
primary_key: Option<String>,
_settings: Vec<(String, String)>,
}
struct FieldAttrs {
name: Option<String>,
primary_key: bool,
skip: bool,
codec: Option<String>,
default: Option<String>,
materialized: Option<String>,
alias_expr: Option<String>,
}
pub fn derive(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let struct_name = &input.ident;
let table_attrs = parse_table_attrs(&input.attrs);
let table_name = table_attrs.name;
let fields = match &input.data {
Data::Struct(data) => match &data.fields {
Fields::Named(fields) => &fields.named,
_ => panic!("#[derive(ClickTable)] only supports structs with named fields"),
},
_ => panic!("#[derive(ClickTable)] only supports structs"),
};
let mut column_structs = vec![];
let mut column_impl = vec![];
let mut column_const_defs = vec![];
let mut column_field_defs = vec![];
let mut column_names = vec![];
let mut insertable_columns = vec![];
let mut primary_key_columns = vec![];
let mut write_row_exprs = vec![];
let mut read_row_exprs = vec![];
let mut static_const_names = Vec::new();
let mut static_const_types = Vec::new();
let mut ddl_columns = vec![];
let mut schema_entries = vec![];
for field in fields {
let field_name = field.ident.as_ref().unwrap();
let field_type = &field.ty;
let attrs = parse_field_attrs(&field.attrs);
if attrs.skip {
continue;
}
let column_name = attrs.name.unwrap_or_else(|| field_name.to_string());
let column_struct_name = format_ident!("{}", to_pascal_case(&column_name));
let column_const_name = format_ident!("{}", field_name.to_string().to_uppercase());
column_names.push(column_name.clone());
if attrs.materialized.is_none() && attrs.alias_expr.is_none() {
insertable_columns.push(column_name.clone());
write_row_exprs.push(quote! {
self.#field_name.write_binary(writer)?;
});
}
if attrs.primary_key {
primary_key_columns.push(column_name.clone());
}
read_row_exprs.push(quote! {
#field_name: clicktype_core::traits::ClickHouseType::read_binary(reader)?,
});
column_structs.push(quote! {
#[derive(Debug, Clone, Copy)]
pub struct #column_struct_name;
});
let qualified_name = format!("{}.{}", table_name, column_name);
let is_pk = attrs.primary_key;
column_impl.push(quote! {
impl clicktype_core::traits::TypedColumn for #column_struct_name {
type Table = super::#struct_name;
type Type = #field_type;
fn name() -> &'static str {
#column_name
}
fn qualified_name() -> &'static str {
#qualified_name
}
fn is_primary_key() -> bool {
#is_pk
}
}
});
column_const_defs.push(quote! {
#field_name: #column_struct_name,
});
column_field_defs.push(quote! {
pub #field_name: #column_struct_name,
});
static_const_names.push(column_const_name);
static_const_types.push(column_struct_name.clone());
ddl_columns.push(field_type);
schema_entries.push(quote! {
(#column_name, <#field_type as clicktype_core::traits::ClickHouseType>::type_name().to_string())
});
}
if let Some(pk_str) = table_attrs.primary_key {
for col_name in pk_str.split(',') {
let trimmed = col_name.trim();
if !trimmed.is_empty() && !primary_key_columns.contains(&trimmed.to_string()) {
primary_key_columns.push(trimmed.to_string());
}
}
}
let columns_mod_name = format_ident!("{}_columns", struct_name.to_string().to_lowercase());
let expanded = quote! {
pub mod #columns_mod_name {
use super::*;
use clicktype_core::traits::TypedColumn;
#(#column_structs)*
#(#column_impl)*
#[derive(Debug, Clone, Copy)]
pub struct Columns {
#(#column_field_defs)*
}
impl Columns {
pub const fn new() -> Self {
Self {
#(#column_const_defs)*
}
}
}
impl Default for Columns {
fn default() -> Self {
Self::new()
}
}
}
impl clicktype_core::traits::ClickTable for #struct_name {
type Columns = #columns_mod_name::Columns;
fn table_name() -> &'static str {
#table_name
}
fn columns() -> Self::Columns {
#columns_mod_name::Columns::default()
}
fn column_names() -> &'static [&'static str] {
&[#(#column_names),*]
}
fn insertable_columns() -> &'static [&'static str] {
&[#(#insertable_columns),*]
}
fn primary_key_columns() -> &'static [&'static str] {
&[#(#primary_key_columns),*]
}
fn create_table_ddl() -> String {
let mut ddl = format!("CREATE TABLE {} (\n", #table_name);
let columns: Vec<(&str, String)> = vec![
#(
(#column_names, <#ddl_columns as clicktype_core::traits::ClickHouseType>::type_name().to_string()),
)*
];
for (i, (name, type_name)) in columns.iter().enumerate() {
if i > 0 {
ddl.push_str(",\n");
}
ddl.push_str(&format!(" {} {}", name, type_name));
}
ddl.push_str("\n) ENGINE = MergeTree()");
let pk_cols: &[&str] = &[#(#primary_key_columns),*];
if !pk_cols.is_empty() {
ddl.push_str(&format!("\nPRIMARY KEY ({})", pk_cols.join(", ")));
}
ddl.push_str("\nORDER BY ");
if !pk_cols.is_empty() {
ddl.push_str(&format!("({})", pk_cols.join(", ")));
} else {
ddl.push_str("tuple()");
}
ddl
}
fn schema() -> Vec<(&'static str, String)> {
vec![
#(#schema_entries),*
]
}
}
impl clicktype_core::traits::ClickInsertable for #struct_name {
fn write_row<W: std::io::Write>(&self, writer: &mut W) -> Result<(), clicktype_core::error::WriteError> {
use clicktype_core::traits::ClickHouseType;
#(#write_row_exprs)*
Ok(())
}
}
impl clicktype_core::traits::ClickReadable for #struct_name {
fn read_row<R: std::io::Read>(reader: &mut R) -> Result<Self, clicktype_core::error::ReadError> {
use clicktype_core::traits::ClickHouseType;
Ok(Self {
#(#read_row_exprs)*
})
}
}
impl #struct_name {
#(pub const #static_const_names: #columns_mod_name::#static_const_types = #columns_mod_name::#static_const_types;)*
}
};
TokenStream::from(expanded)
}
fn parse_table_attrs(attrs: &[Attribute]) -> TableAttrs {
let mut name = None;
let mut engine = None;
let mut partition_by = None;
let mut order_by = None;
let mut primary_key = None;
let settings = Vec::new();
for attr in attrs {
if !attr.path().is_ident("click_table") {
continue;
}
let _ = attr.parse_nested_meta(|meta| {
if meta.path.is_ident("name") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
name = Some(s.value());
}
} else if meta.path.is_ident("engine") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
engine = Some(s.value());
}
} else if meta.path.is_ident("partition_by") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
partition_by = Some(s.value());
}
} else if meta.path.is_ident("order_by") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
order_by = Some(s.value());
}
} else if meta.path.is_ident("primary_key") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
primary_key = Some(s.value());
}
}
Ok(())
});
}
TableAttrs {
name: name.expect("click_table attribute requires 'name' parameter"),
_engine: engine,
_partition_by: partition_by,
_order_by: order_by,
primary_key,
_settings: settings,
}
}
fn parse_field_attrs(attrs: &[Attribute]) -> FieldAttrs {
let mut field_attrs = FieldAttrs {
name: None,
primary_key: false,
skip: false,
codec: None,
default: None,
materialized: None,
alias_expr: None,
};
for attr in attrs {
if !attr.path().is_ident("click_column") {
continue;
}
let _ = attr.parse_nested_meta(|meta| {
if meta.path.is_ident("primary_key") {
field_attrs.primary_key = true;
} else if meta.path.is_ident("skip") {
field_attrs.skip = true;
} else if meta.path.is_ident("rename") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
field_attrs.name = Some(s.value());
}
} else if meta.path.is_ident("codec") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
field_attrs.codec = Some(s.value());
}
} else if meta.path.is_ident("default") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
field_attrs.default = Some(s.value());
}
} else if meta.path.is_ident("materialized") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
field_attrs.materialized = Some(s.value());
}
} else if meta.path.is_ident("alias") {
let value: Lit = meta.value()?.parse()?;
if let Lit::Str(s) = value {
field_attrs.alias_expr = Some(s.value());
}
}
Ok(())
});
}
field_attrs
}
fn to_pascal_case(s: &str) -> String {
s.split('_')
.map(|word| {
let mut chars = word.chars();
match chars.next() {
None => String::new(),
Some(first) => {
first.to_uppercase().chain(chars).collect()
}
}
})
.collect()
}