spyglass_cli/
codegen.rs

1use crate::graphql_types::*;
2use anyhow::Error;
3use proc_macro2::TokenStream;
4use quote::{format_ident, quote};
5
6pub fn generate_schema_code(schema: Schema) -> Result<Vec<(TokenStream, String)>, Error> {
7    let types = schema.types;
8    let type_definitions: Vec<(proc_macro2::TokenStream, String)> = types
9        .iter()
10        .filter_map(|type_definition| {
11            if type_definition.name == "Query" {
12                None
13            } else {
14                let module = generate_graphql_helpers(type_definition.clone()).unwrap();
15                Some((module, type_definition.snake_case_name()))
16            }
17        })
18        .collect();
19
20    Ok(type_definitions)
21}
22
23fn generate_graphql_helpers(entity: TypeDefinition) -> Result<TokenStream, Error> {
24    let TypeDefinition { name, fields } = entity;
25
26    let module_name = format_ident!("{}", &name);
27
28    let struct_definition = generate_struct(&fields);
29
30    let table_name = format!("{}", &name);
31    let create_function_params = generate_create_params(&fields);
32    let create_function_set_statements = generate_create_set_statements(&fields);
33    let create_function = quote!(
34        pub fn create(tables: &mut Tables, id: &String, #create_function_params) {
35            let row = tables.create_row(#table_name, id);
36            #create_function_set_statements
37        }
38    );
39
40    let update_functions = generate_update_functions(&table_name, &fields);
41
42    let module_definition = quote!(
43        pub mod #module_name {
44            use substreams_entity_change::tables::Tables;
45
46            #struct_definition
47
48            #create_function
49
50            #update_functions
51        }
52    );
53
54    Ok(quote!(
55        #module_definition
56    )
57    .into())
58}
59
60fn generate_struct(fields: &Vec<Field>) -> proc_macro2::TokenStream {
61    let field_definitions: Vec<proc_macro2::TokenStream> = fields
62        .iter()
63        .map(|field| {
64            let field_name = format_ident!("r#{}", &field.name);
65            let field_type = field.rust_type();
66            quote!(pub #field_name: #field_type)
67        })
68        .collect();
69
70    quote!(
71        pub struct Model {
72            #(#field_definitions),*
73        }
74    )
75}
76
77fn generate_create_params(fields: &Vec<Field>) -> proc_macro2::TokenStream {
78    let field_types: Vec<proc_macro2::TokenStream> = fields
79        .iter()
80        .filter_map(|field| {
81            if field.is_id() || !field.required() || field.derived_from() {
82                None
83            } else {
84                let field_type = field.rust_type();
85                let field_name = format_ident!("r#{}", field.snake_case_name());
86                Some(quote!(#field_name: #field_type))
87            }
88        })
89        .collect();
90
91    quote! {
92        #(#field_types),*
93    }
94}
95
96fn generate_set_statement(field: &Field) -> proc_macro2::TokenStream {
97    let key = format!("r#{}", &field.name());
98    let value = format_ident!("r#{}", &field.snake_case_name());
99    match field.field_type() {
100        GraphQlType::BigInt { .. } => {
101            quote!(
102                row.set_bigint(#key, #value);
103            )
104        }
105        _ => {
106            quote!(
107                row.set(#key, #value);
108            )
109        }
110    }
111}
112
113pub fn generate_update_functions(
114    table_name: &String,
115    fields: &Vec<Field>,
116) -> proc_macro2::TokenStream {
117    let update_functions: Vec<proc_macro2::TokenStream> = fields
118        .iter()
119        .filter_map(|field| {
120            if field.is_id() || field.derived_from() {
121                None
122            } else {
123                let table_name = format!("{}", table_name);
124                let argument_name = format_ident!("r#{}", field.snake_case_name());
125                let field_type = field.rust_type();
126                let update_function_name = format_ident!("update_{}", field.snake_case_name());
127
128                let set_statement = generate_set_statement(field);
129
130                Some(quote!(
131                    pub fn #update_function_name(tables: &mut Tables, id: &String, #argument_name: #field_type) {
132                        let row = tables.update_row(#table_name, id);
133                        #set_statement
134                    }
135                ))
136            }
137        })
138        .collect();
139
140    quote! {
141        #(#update_functions)*
142    }
143}
144fn generate_create_set_statements(fields: &Vec<Field>) -> proc_macro2::TokenStream {
145    let set_statements: Vec<proc_macro2::TokenStream> = fields
146        .iter()
147        .filter_map(|field| {
148            if field.is_id() || !field.required() || field.derived_from() {
149                None
150            } else {
151                Some(generate_set_statement(&field))
152            }
153        })
154        .collect();
155
156    quote! {
157        #(#set_statements)*
158    }
159}