typed_store_derive/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::collections::{BTreeMap, HashSet};
6
7use itertools::Itertools;
8use proc_macro::TokenStream;
9use proc_macro2::Ident;
10use quote::quote;
11use syn::{
12    AngleBracketedGenericArguments, Attribute, Generics, ItemStruct, Lit, Meta, PathArguments,
13    Type::{self},
14    parse_macro_input,
15};
16
17// This is used as default when none is specified
18const DEFAULT_DB_OPTIONS_CUSTOM_FN: &str = "typed_store::rocks::default_db_options";
19// Custom function which returns the option and overrides the defaults for this
20// table
21const DB_OPTIONS_CUSTOM_FUNCTION: &str = "default_options_override_fn";
22// Use a different name for the column than the identifier
23const DB_OPTIONS_RENAME: &str = "rename";
24// Deprecate a column family
25const DB_OPTIONS_DEPRECATE: &str = "deprecated";
26
27/// Options can either be simplified form or
28enum GeneralTableOptions {
29    OverrideFunction(String),
30}
31
32impl Default for GeneralTableOptions {
33    fn default() -> Self {
34        Self::OverrideFunction(DEFAULT_DB_OPTIONS_CUSTOM_FN.to_owned())
35    }
36}
37
38// Extracts the field names, field types, inner types (K,V in {map_type_name}<K,
39// V>), and the options attrs
40fn extract_struct_info(
41    input: ItemStruct,
42    allowed_map_type_names: HashSet<String>,
43) -> ExtractedStructInfo {
44    // There must only be one map type used for all entries
45    let allowed_strs: Vec<_> = allowed_map_type_names
46        .iter()
47        .map(|s| format!("{s}<K, V>"))
48        .collect();
49    let allowed_strs = allowed_strs.join(" or ");
50    let mut deprecated_cfs = vec![];
51
52    let info = input.fields.iter().map(|f| {
53        let attrs: BTreeMap<_, _> = f
54            .attrs
55            .iter()
56            .filter(|a| {
57                a.path.is_ident(DB_OPTIONS_CUSTOM_FUNCTION)
58                    || a.path.is_ident(DB_OPTIONS_RENAME)
59                    || a.path.is_ident(DB_OPTIONS_DEPRECATE)
60            })
61            .map(|a| (a.path.get_ident().unwrap().to_string(), a))
62            .collect();
63
64        let options = if let Some(options) = attrs.get(DB_OPTIONS_CUSTOM_FUNCTION) {
65            GeneralTableOptions::OverrideFunction(get_options_override_function(options).unwrap())
66        } else {
67            GeneralTableOptions::default()
68        };
69
70        let ty = &f.ty;
71        if let Type::Path(p) = ty {
72            let type_info = &p.path.segments.first().unwrap();
73            let inner_type =
74                if let PathArguments::AngleBracketed(angle_bracket_type) = &type_info.arguments {
75                    angle_bracket_type.clone()
76                } else {
77                    panic!("All struct members must be of type {allowed_strs}");
78                };
79
80            let type_str = format!("{}", &type_info.ident);
81            // Rough way to check that this is map_type_name
82            if allowed_map_type_names.contains(&type_str) {
83                let field_name = f.ident.as_ref().unwrap().clone();
84                let cf_name = if let Some(rename) = attrs.get(DB_OPTIONS_RENAME) {
85                    match rename.parse_meta().expect("Cannot parse meta of attribute") {
86                        Meta::NameValue(val) => {
87                            if let Lit::Str(s) = val.lit {
88                                // convert to ident
89                                s.parse().expect("Rename value must be identifier")
90                            } else {
91                                panic!("Expected string value for rename")
92                            }
93                        }
94                        _ => panic!("Expected string value for rename"),
95                    }
96                } else {
97                    field_name.clone()
98                };
99                if attrs.contains_key(DB_OPTIONS_DEPRECATE) {
100                    deprecated_cfs.push(field_name.clone());
101                }
102
103                return ((field_name, cf_name, type_str), (inner_type, options));
104            } else {
105                panic!("All struct members must be of type {allowed_strs}");
106            }
107        }
108        panic!("All struct members must be of type {allowed_strs}");
109    });
110
111    let (field_info, inner_types_with_opts): (Vec<_>, Vec<_>) = info.unzip();
112    let (field_names, cf_names, simple_field_type_names): (Vec<_>, Vec<_>, Vec<_>) =
113        field_info.into_iter().multiunzip();
114
115    // Check for homogeneous types
116    if let Some(first) = simple_field_type_names.first() {
117        simple_field_type_names.iter().for_each(|q| {
118            if q != first {
119                panic!("All struct members must be of same type");
120            }
121        })
122    } else {
123        panic!("Cannot derive on empty struct");
124    };
125
126    let (inner_types, options): (Vec<_>, Vec<_>) = inner_types_with_opts.into_iter().unzip();
127
128    ExtractedStructInfo {
129        field_names,
130        cf_names,
131        inner_types,
132        derived_table_options: options,
133        simple_field_type_name_str: simple_field_type_names.first().unwrap().clone(),
134        deprecated_cfs,
135    }
136}
137
138/// Extracts the table options override function
139/// The function must take no args and return Options
140fn get_options_override_function(attr: &Attribute) -> syn::Result<String> {
141    let meta = attr.parse_meta()?;
142
143    let val = match meta.clone() {
144        Meta::NameValue(val) => val,
145        _ => {
146            return Err(syn::Error::new_spanned(
147                meta,
148                format!(
149                    "Expected function name in format `#[{DB_OPTIONS_CUSTOM_FUNCTION} = {{function_name}}]`"
150                ),
151            ));
152        }
153    };
154
155    if !val.path.is_ident(DB_OPTIONS_CUSTOM_FUNCTION) {
156        return Err(syn::Error::new_spanned(
157            meta,
158            format!(
159                "Expected function name in format `#[{DB_OPTIONS_CUSTOM_FUNCTION} = {{function_name}}]`"
160            ),
161        ));
162    }
163
164    let fn_name = match val.lit {
165        Lit::Str(fn_name) => fn_name,
166        _ => {
167            return Err(syn::Error::new_spanned(
168                meta,
169                format!(
170                    "Expected function name in format `#[{DB_OPTIONS_CUSTOM_FUNCTION} = {{function_name}}]`"
171                ),
172            ));
173        }
174    };
175    Ok(fn_name.value())
176}
177
178fn extract_generics_names(generics: &Generics) -> Vec<Ident> {
179    generics
180        .params
181        .iter()
182        .map(|g| match g {
183            syn::GenericParam::Type(t) => t.ident.clone(),
184            _ => panic!("Unsupported generic type"),
185        })
186        .collect()
187}
188
189struct ExtractedStructInfo {
190    field_names: Vec<Ident>,
191    cf_names: Vec<Ident>,
192    inner_types: Vec<AngleBracketedGenericArguments>,
193    derived_table_options: Vec<GeneralTableOptions>,
194    simple_field_type_name_str: String,
195    deprecated_cfs: Vec<Ident>,
196}
197
198#[proc_macro_derive(DBMapUtils, attributes(default_options_override_fn, rename))]
199pub fn derive_dbmap_utils_general(input: TokenStream) -> TokenStream {
200    let input = parse_macro_input!(input as ItemStruct);
201    let name = &input.ident;
202    let generics = &input.generics;
203    let generics_names = extract_generics_names(generics);
204
205    let allowed_types_with_post_process_fn: BTreeMap<_, _> =
206        [("SallyColumn", ""), ("DBMap", "")].into_iter().collect();
207    let allowed_strs = allowed_types_with_post_process_fn
208        .keys()
209        .map(|s| s.to_string())
210        .collect();
211
212    // TODO: use `parse_quote` over `parse()`
213    let ExtractedStructInfo {
214        field_names,
215        cf_names,
216        inner_types,
217        derived_table_options,
218        simple_field_type_name_str,
219        deprecated_cfs,
220    } = extract_struct_info(input.clone(), allowed_strs);
221
222    let (key_names, value_names): (Vec<_>, Vec<_>) = inner_types
223        .iter()
224        .map(|q| (q.args.first().unwrap(), q.args.last().unwrap()))
225        .unzip();
226
227    // This is the actual name of the type which was found
228    let post_process_fn_str = allowed_types_with_post_process_fn
229        .get(&simple_field_type_name_str.as_str())
230        .unwrap();
231    let post_process_fn: proc_macro2::TokenStream = post_process_fn_str.parse().unwrap();
232
233    let default_options_override_fn_names: Vec<proc_macro2::TokenStream> = derived_table_options
234        .iter()
235        .map(|q| {
236            let GeneralTableOptions::OverrideFunction(fn_name) = q;
237            fn_name.parse().unwrap()
238        })
239        .collect();
240
241    let generics_bounds =
242        "std::fmt::Debug + serde::Serialize + for<'de> serde::de::Deserialize<'de>";
243    let generics_bounds_token: proc_macro2::TokenStream = generics_bounds.parse().unwrap();
244
245    let config_struct_name_str = format!("{name}Configurator");
246    let config_struct_name: proc_macro2::TokenStream = config_struct_name_str.parse().unwrap();
247
248    let intermediate_db_map_struct_name_str = format!("{name}IntermediateDBMapStructPrimary");
249    let intermediate_db_map_struct_name: proc_macro2::TokenStream =
250        intermediate_db_map_struct_name_str.parse().unwrap();
251
252    let secondary_db_map_struct_name_str = format!("{name}ReadOnly");
253    let secondary_db_map_struct_name: proc_macro2::TokenStream =
254        secondary_db_map_struct_name_str.parse().unwrap();
255
256    TokenStream::from(quote! {
257
258        // <----------- This section generates the configurator struct -------------->
259
260        /// Create config structs for configuring DBMap tables
261        pub struct #config_struct_name {
262            #(
263                pub #field_names : typed_store::rocks::DBOptions,
264            )*
265        }
266
267        impl #config_struct_name {
268            /// Initialize to defaults
269            pub fn init() -> Self {
270                Self {
271                    #(
272                        #field_names : typed_store::rocks::default_db_options(),
273                    )*
274                }
275            }
276
277            /// Build a config
278            pub fn build(&self) -> typed_store::rocks::DBMapTableConfigMap {
279                typed_store::rocks::DBMapTableConfigMap::new([
280                    #(
281                        (stringify!(#field_names).to_owned(), self.#field_names.clone()),
282                    )*
283                ].into_iter().collect())
284            }
285        }
286
287        impl <
288                #(
289                    #generics_names: #generics_bounds_token,
290                )*
291            > #name #generics {
292
293                pub fn configurator() -> #config_struct_name {
294                    #config_struct_name::init()
295                }
296        }
297
298        // <----------- This section generates the core open logic for opening DBMaps -------------->
299
300        /// Create an intermediate struct used to open the DBMap tables in primary mode
301        /// This is only used internally
302        struct #intermediate_db_map_struct_name #generics {
303                #(
304                    pub #field_names : DBMap #inner_types,
305                )*
306        }
307
308
309        impl <
310                #(
311                    #generics_names: #generics_bounds_token,
312                )*
313            > #intermediate_db_map_struct_name #generics {
314            /// Opens a set of tables in read-write mode
315            /// If as_secondary_with_path is set, the DB is opened in read only mode with the path specified
316            pub fn open_tables_impl(
317                path: std::path::PathBuf,
318                as_secondary_with_path: Option<std::path::PathBuf>,
319                is_transaction: bool,
320                metric_conf: typed_store::rocks::MetricConf,
321                global_db_options_override: Option<typed_store::rocksdb::Options>,
322                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>,
323                remove_deprecated_tables: bool,
324            ) -> Self {
325                let path = &path;
326                let default_cf_opt = if let Some(opt) = global_db_options_override.as_ref() {
327                    typed_store::rocks::DBOptions {
328                        options: opt.clone(),
329                        rw_options: typed_store::rocks::default_db_options().rw_options,
330                    }
331                } else {
332                    typed_store::rocks::default_db_options()
333                };
334                let (db, rwopt_cfs) = {
335                    let opt_cfs = match tables_db_options_override {
336                        None => [
337                            #(
338                                (stringify!(#cf_names).to_owned(), #default_options_override_fn_names()),
339                            )*
340                        ],
341                        Some(o) => [
342                            #(
343                                (stringify!(#cf_names).to_owned(), o.to_map().get(stringify!(#cf_names)).unwrap_or(&default_cf_opt).clone()),
344                            )*
345                        ]
346                    };
347                    // Safe to call unwrap because we will have at least one field_name entry in the struct
348                    let rwopt_cfs: std::collections::HashMap<String, typed_store::rocks::ReadWriteOptions> = opt_cfs.iter().map(|q| (q.0.as_str().to_string(), q.1.rw_options.clone())).collect();
349                    let opt_cfs: Vec<_> = opt_cfs.iter().map(|q| (q.0.as_str(), q.1.options.clone())).collect();
350                    let db = match (as_secondary_with_path.clone(), is_transaction) {
351                        (Some(p), _) => typed_store::rocks::open_cf_opts_secondary(path, Some(&p), global_db_options_override, metric_conf, &opt_cfs),
352                        (_, true) => typed_store::rocks::open_cf_opts_transactional(path, global_db_options_override, metric_conf, &opt_cfs),
353                        _ => typed_store::rocks::open_cf_opts(path, global_db_options_override, metric_conf, &opt_cfs)
354                    };
355                    db.map(|d| (d, rwopt_cfs))
356                }.expect(&format!("Cannot open DB at {:?}", path));
357                let deprecated_tables = vec![#(stringify!(#deprecated_cfs),)*];
358                let (
359                        #(
360                            #field_names
361                        ),*
362                ) = (#(
363                        DBMap::#inner_types::reopen(&db, Some(stringify!(#cf_names)), rwopt_cfs.get(stringify!(#cf_names)).unwrap_or(&typed_store::rocks::ReadWriteOptions::default()), remove_deprecated_tables && deprecated_tables.contains(&stringify!(#cf_names))).expect(&format!("Cannot open {} CF.", stringify!(#cf_names))[..])
364                    ),*);
365
366                if as_secondary_with_path.is_none() && remove_deprecated_tables {
367                    #(
368                        db.drop_cf(stringify!(#deprecated_cfs)).expect("failed to drop a deprecated cf");
369                    )*
370                }
371                Self {
372                    #(
373                        #field_names,
374                    )*
375                }
376            }
377        }
378
379
380        // <----------- This section generates the read-write open logic and other common utils -------------->
381
382        impl <
383                #(
384                    #generics_names: #generics_bounds_token,
385                )*
386            > #name #generics {
387            /// Opens a set of tables in read-write mode
388            /// Only one process is allowed to do this at a time
389            /// `global_db_options_override` apply to the whole DB
390            /// `tables_db_options_override` apply to each table. If `None`, the attributes from `default_options_override_fn` are used if any
391            #[expect(unused_parens)]
392            pub fn open_tables_read_write(
393                path: std::path::PathBuf,
394                metric_conf: typed_store::rocks::MetricConf,
395                global_db_options_override: Option<typed_store::rocksdb::Options>,
396                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>
397            ) -> Self {
398                let inner = #intermediate_db_map_struct_name::open_tables_impl(path, None, false, metric_conf, global_db_options_override, tables_db_options_override, false);
399                Self {
400                    #(
401                        #field_names: #post_process_fn(inner.#field_names),
402                    )*
403                }
404            }
405
406            #[expect(unused_parens)]
407            pub fn open_tables_read_write_with_deprecation_option(
408                path: std::path::PathBuf,
409                metric_conf: typed_store::rocks::MetricConf,
410                global_db_options_override: Option<typed_store::rocksdb::Options>,
411                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>,
412                remove_deprecated_tables: bool,
413            ) -> Self {
414                let inner = #intermediate_db_map_struct_name::open_tables_impl(path, None, false, metric_conf, global_db_options_override, tables_db_options_override, remove_deprecated_tables);
415                Self {
416                    #(
417                        #field_names: #post_process_fn(inner.#field_names),
418                    )*
419                }
420            }
421
422            /// Opens a set of tables in transactional read-write mode
423            /// Only one process is allowed to do this at a time
424            /// `global_db_options_override` apply to the whole DB
425            /// `tables_db_options_override` apply to each table. If `None`, the attributes from `default_options_override_fn` are used if any
426            #[expect(unused_parens)]
427            pub fn open_tables_transactional(
428                path: std::path::PathBuf,
429                metric_conf: typed_store::rocks::MetricConf,
430                global_db_options_override: Option<typed_store::rocksdb::Options>,
431                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>
432            ) -> Self {
433                let inner = #intermediate_db_map_struct_name::open_tables_impl(path, None, true, metric_conf, global_db_options_override, tables_db_options_override, false);
434                Self {
435                    #(
436                        #field_names: #post_process_fn(inner.#field_names),
437                    )*
438                }
439            }
440
441            /// Returns a list of the tables name and type pairs
442            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
443                vec![#(
444                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
445                )*].into_iter().collect()
446            }
447
448            /// This opens the DB in read only mode and returns a struct which exposes debug features
449            pub fn get_read_only_handle (
450                primary_path: std::path::PathBuf,
451                with_secondary_path: Option<std::path::PathBuf>,
452                global_db_options_override: Option<typed_store::rocksdb::Options>,
453                metric_conf: typed_store::rocks::MetricConf,
454                ) -> #secondary_db_map_struct_name #generics {
455                #secondary_db_map_struct_name::open_tables_read_only(primary_path, with_secondary_path, metric_conf, global_db_options_override)
456            }
457        }
458
459
460        // <----------- This section generates the features that use read-only open logic -------------->
461        /// Create an intermediate struct used to open the DBMap tables in secondary mode
462        /// This is only used internally
463        pub struct #secondary_db_map_struct_name #generics {
464            #(
465                pub #field_names : DBMap #inner_types,
466            )*
467        }
468
469        impl <
470                #(
471                    #generics_names: #generics_bounds_token,
472                )*
473            > #secondary_db_map_struct_name #generics {
474            /// Open in read only mode. No limitation on number of processes to do this
475            pub fn open_tables_read_only(
476                primary_path: std::path::PathBuf,
477                with_secondary_path: Option<std::path::PathBuf>,
478                metric_conf: typed_store::rocks::MetricConf,
479                global_db_options_override: Option<typed_store::rocksdb::Options>,
480            ) -> Self {
481                let inner = match with_secondary_path {
482                    Some(q) => #intermediate_db_map_struct_name::open_tables_impl(primary_path, Some(q), false, metric_conf, global_db_options_override, None, false),
483                    None => {
484                        let p: std::path::PathBuf = tempfile::tempdir()
485                        .expect("Failed to open temporary directory")
486                        .into_path();
487                        #intermediate_db_map_struct_name::open_tables_impl(primary_path, Some(p), false, metric_conf, global_db_options_override, None, false)
488                    }
489                };
490                Self {
491                    #(
492                        #field_names: inner.#field_names,
493                    )*
494                }
495            }
496
497            fn cf_name_to_table_name(cf_name: &str) -> eyre::Result<&'static str> {
498                Ok(match cf_name {
499                    #(
500                        stringify!(#cf_names) => stringify!(#field_names),
501                    )*
502                    _ => eyre::bail!("No such cf name: {}", cf_name),
503                })
504            }
505
506            /// Dump all key-value pairs in the page at the given table name
507            /// Tables must be opened in read only mode using `open_tables_read_only`
508            pub fn dump(&self, cf_name: &str, page_size: u16, page_number: usize) -> eyre::Result<std::collections::BTreeMap<String, String>> {
509                let table_name = Self::cf_name_to_table_name(cf_name)?;
510
511                Ok(match table_name {
512                    #(
513                        stringify!(#field_names) => {
514                            typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
515                            typed_store::traits::Map::unbounded_iter(&self.#field_names)
516                                .skip((page_number * (page_size) as usize))
517                                .take(page_size as usize)
518                                .map(|(k, v)| (format!("{:?}", k), format!("{:?}", v)))
519                                .collect::<std::collections::BTreeMap<_, _>>()
520                        }
521                    )*
522
523                    _ => eyre::bail!("No such table name: {}", table_name),
524                })
525            }
526
527            /// Get key value sizes from the db
528            /// Tables must be opened in read only mode using `open_tables_read_only`
529            pub fn table_summary(&self, table_name: &str) -> eyre::Result<typed_store::traits::TableSummary> {
530                let mut count = 0;
531                let mut key_bytes = 0;
532                let mut value_bytes = 0;
533                match table_name {
534                    #(
535                        stringify!(#field_names) => {
536                            typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
537                            self.#field_names.table_summary()
538                        }
539                    )*
540
541                    _ => eyre::bail!("No such table name: {}", table_name),
542                }
543            }
544
545            /// Count the keys in this table
546            /// Tables must be opened in read only mode using `open_tables_read_only`
547            pub fn count_keys(&self, table_name: &str) -> eyre::Result<usize> {
548                Ok(match table_name {
549                    #(
550                        stringify!(#field_names) => {
551                            typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
552                            typed_store::traits::Map::unbounded_iter(&self.#field_names).count()
553                        }
554                    )*
555
556                    _ => eyre::bail!("No such table name: {}", table_name),
557                })
558            }
559
560            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
561                vec![#(
562                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
563                )*].into_iter().collect()
564            }
565
566            /// Try catch up with primary for all tables. This can be a slow operation
567            /// Tables must be opened in read only mode using `open_tables_read_only`
568            pub fn try_catch_up_with_primary_all(&self) -> eyre::Result<()> {
569                #(
570                    typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
571                )*
572                Ok(())
573            }
574        }
575
576        impl <
577                #(
578                    #generics_names: #generics_bounds_token,
579                )*
580            > TypedStoreDebug for #secondary_db_map_struct_name #generics {
581                fn dump_table(
582                    &self,
583                    table_name: String,
584                    page_size: u16,
585                    page_number: usize,
586                ) -> eyre::Result<std::collections::BTreeMap<String, String>> {
587                    self.dump(table_name.as_str(), page_size, page_number)
588                }
589
590                fn primary_db_name(&self) -> String {
591                    stringify!(#name).to_owned()
592                }
593
594                fn describe_all_tables(&self) -> std::collections::BTreeMap<String, (String, String)> {
595                    Self::describe_tables()
596                }
597
598                fn count_table_keys(&self, table_name: String) -> eyre::Result<usize> {
599                    self.count_keys(table_name.as_str())
600                }
601
602                fn table_summary(&self, table_name: String) -> eyre::Result<TableSummary> {
603                    self.table_summary(table_name.as_str())
604                }
605
606
607        }
608
609    })
610}
611
612#[proc_macro_derive(SallyDB, attributes(default_options_override_fn))]
613pub fn derive_sallydb_general(input: TokenStream) -> TokenStream {
614    // log_syntax!("here");
615    let input = parse_macro_input!(input as ItemStruct);
616    let name = &input.ident;
617    let generics = &input.generics;
618    let generics_names = extract_generics_names(generics);
619
620    let allowed_types_with_post_process_fn: BTreeMap<_, _> =
621        [("SallyColumn", "")].into_iter().collect();
622    let allowed_strs = allowed_types_with_post_process_fn
623        .keys()
624        .map(|s| s.to_string())
625        .collect();
626
627    // TODO: use `parse_quote` over `parse()`
628    // TODO: Eventually this should return a Vec<Vec<GeneralTableOptions>> to
629    // capture default table options for each column type i.e. RockDB, TestDB, etc
630    let ExtractedStructInfo {
631        field_names,
632        inner_types,
633        derived_table_options,
634        simple_field_type_name_str,
635        ..
636    } = extract_struct_info(input.clone(), allowed_strs);
637
638    let (key_names, value_names): (Vec<_>, Vec<_>) = inner_types
639        .iter()
640        .map(|q| (q.args.first().unwrap(), q.args.last().unwrap()))
641        .unzip();
642
643    // This is the actual name of the type which was found
644    let post_process_fn_str = allowed_types_with_post_process_fn
645        .get(&simple_field_type_name_str.as_str())
646        .unwrap();
647    let post_process_fn: proc_macro2::TokenStream = post_process_fn_str.parse().unwrap();
648
649    let default_options_override_fn_names: Vec<proc_macro2::TokenStream> = derived_table_options
650        .iter()
651        .map(|q| {
652            let GeneralTableOptions::OverrideFunction(fn_name) = q;
653            fn_name.parse().unwrap()
654        })
655        .collect();
656
657    let generics_bounds =
658        "std::fmt::Debug + serde::Serialize + for<'de> serde::de::Deserialize<'de>";
659    let generics_bounds_token: proc_macro2::TokenStream = generics_bounds.parse().unwrap();
660
661    let config_struct_name_str = format!("{name}SallyConfigurator");
662    let sally_config_struct_name: proc_macro2::TokenStream =
663        config_struct_name_str.parse().unwrap();
664
665    let intermediate_db_map_struct_name_str = format!("{name}Primary");
666    let intermediate_db_map_struct_name: proc_macro2::TokenStream =
667        intermediate_db_map_struct_name_str.parse().unwrap();
668
669    let secondary_db_map_struct_name_str = format!("{name}ReadOnly");
670    let secondary_db_map_struct_name: proc_macro2::TokenStream =
671        secondary_db_map_struct_name_str.parse().unwrap();
672
673    TokenStream::from(quote! {
674
675        // <----------- This section generates the configurator struct -------------->
676
677        /// Create config structs for configuring SallyColumns
678        pub struct #sally_config_struct_name {
679            #(
680                pub #field_names : typed_store::sally::SallyColumnOptions,
681            )*
682        }
683
684        impl #sally_config_struct_name {
685            /// Initialize to defaults
686            pub fn init() -> Self {
687                Self {
688                    #(
689                        #field_names : typed_store::sally::default_column_options(),
690                    )*
691                }
692            }
693
694            /// Build a config
695            pub fn build(&self) -> typed_store::sally::SallyDBConfigMap {
696                typed_store::sally::SallyDBConfigMap::new([
697                    #(
698                        (stringify!(#field_names).to_owned(), self.#field_names.clone()),
699                    )*
700                ].into_iter().collect())
701            }
702        }
703
704
705        impl <
706                #(
707                    #generics_names: #generics_bounds_token,
708                )*
709            > #name #generics {
710
711                pub fn configurator() -> #sally_config_struct_name {
712                    #sally_config_struct_name::init()
713                }
714        }
715
716
717        // <----------- This section generates the core open logic for opening sally columns -------------->
718
719        /// Create an intermediate struct used to open the DBMap tables in primary mode
720        /// This is only used internally
721        struct #intermediate_db_map_struct_name #generics {
722                #(
723                    pub #field_names : SallyColumn #inner_types,
724                )*
725        }
726
727
728        impl <
729                #(
730                    #generics_names: #generics_bounds_token,
731                )*
732            > #intermediate_db_map_struct_name #generics {
733            /// Opens a set of tables in read-write mode
734            /// If as_secondary_with_path is set, the DB is opened in read only mode with the path specified
735            pub fn init(db_options: typed_store::sally::SallyDBOptions) -> Self {
736                match db_options {
737                    typed_store::sally::SallyDBOptions::TestDB => {
738                        let (
739                            #(
740                                #field_names
741                            ),*
742                        ) = (#(
743                            SallyColumn::TestDB((typed_store::test_db::TestDB::#inner_types::open(), typed_store::sally::SallyConfig::default()))
744                            ),*);
745
746                        Self {
747                            #(
748                                #field_names,
749                            )*
750                        }
751                    },
752                    typed_store::sally::SallyDBOptions::RocksDB((path, metric_conf, access_type, global_db_options_override, tables_db_options_override)) => {
753                        let path = &path;
754                        let (db, rwopt_cfs) = {
755                            let opt_cfs = match tables_db_options_override {
756                                None => [
757                                    #(
758                                        (stringify!(#field_names).to_owned(), #default_options_override_fn_names().clone()),
759                                    )*
760                                ],
761                                Some(o) => [
762                                    #(
763                                        (stringify!(#field_names).to_owned(), o.to_map().get(stringify!(#field_names)).unwrap().clone()),
764                                    )*
765                                ]
766                            };
767                            // Safe to call unwrap because we will have at least one field_name entry in the struct
768                            let rwopt_cfs: std::collections::HashMap<String, typed_store::rocks::ReadWriteOptions> = opt_cfs.iter().map(|q| (q.0.as_str().to_string(), q.1.rw_options.clone())).collect();
769                            let opt_cfs: Vec<_> = opt_cfs.iter().map(|q| (q.0.as_str(), q.1.options.clone())).collect();
770                            let db = match access_type {
771                                RocksDBAccessType::Secondary(Some(p)) => typed_store::rocks::open_cf_opts_secondary(path, Some(&p), global_db_options_override, metric_conf, &opt_cfs),
772                                _ => typed_store::rocks::open_cf_opts(path, global_db_options_override, metric_conf, &opt_cfs)
773                            };
774                            db.map(|d| (d, rwopt_cfs))
775                        }.expect(&format!("Cannot open DB at {:?}", path));
776                        let (
777                            #(
778                                #field_names
779                            ),*
780                        ) = (#(
781                            SallyColumn::RocksDB((DBMap::#inner_types::reopen(&db, Some(stringify!(#field_names)), rwopt_cfs.get(stringify!(#field_names)).unwrap_or(&typed_store::rocks::ReadWriteOptions::default()), false).expect(&format!("Cannot open {} CF.", stringify!(#field_names))[..]), typed_store::sally::SallyConfig::default()))
782                            ),*);
783
784                        Self {
785                            #(
786                                #field_names,
787                            )*
788                        }
789                    }
790                }
791            }
792        }
793
794
795        // <----------- This section generates the read-write open logic and other common utils -------------->
796        impl <
797                #(
798                    #generics_names: #generics_bounds_token,
799                )*
800            > #name #generics {
801            /// Opens a set of tables in read-write mode
802            /// Only one process is allowed to do this at a time
803            /// `global_db_options_override` apply to the whole DB
804            /// `tables_db_options_override` apply to each table. If `None`, the attributes from `default_options_override_fn` are used if any
805            #[expect(unused_parens)]
806            pub fn init(
807                db_options: typed_store::sally::SallyDBOptions
808            ) -> Self {
809                let inner = #intermediate_db_map_struct_name::init(db_options);
810                Self {
811                    #(
812                        #field_names: #post_process_fn(inner.#field_names),
813                    )*
814                }
815            }
816
817            /// Returns a list of the tables name and type pairs
818            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
819                vec![#(
820                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
821                )*].into_iter().collect()
822            }
823
824            /// This opens the DB in read only mode and returns a struct which exposes debug features
825            pub fn get_read_only_handle (
826                db_options: typed_store::sally::SallyReadOnlyDBOptions
827                ) -> #secondary_db_map_struct_name #generics {
828                #secondary_db_map_struct_name::init_read_only(db_options)
829            }
830        }
831
832        // <----------- This section generates the features that use read-only open logic -------------->
833        /// Create an intermediate struct used to open the DBMap tables in secondary mode
834        /// This is only used internally
835        pub struct #secondary_db_map_struct_name #generics {
836            #(
837                pub #field_names : SallyColumn #inner_types,
838            )*
839        }
840
841        impl <
842                #(
843                    #generics_names: #generics_bounds_token,
844                )*
845            > #secondary_db_map_struct_name #generics {
846            /// Open in read only mode. No limitation on number of processes to do this
847            pub fn init_read_only(
848                db_options: typed_store::sally::SallyReadOnlyDBOptions,
849            ) -> Self {
850                match db_options {
851                    typed_store::sally::SallyReadOnlyDBOptions::TestDB => {
852                        let inner = #intermediate_db_map_struct_name::init(SallyDBOptions::TestDB);
853                        Self {
854                            #(
855                                #field_names: inner.#field_names,
856                            )*
857                        }
858                    },
859                    typed_store::sally::SallyReadOnlyDBOptions::RocksDB(b) => {
860                        let inner = match b.2 {
861                            Some(q) => #intermediate_db_map_struct_name::init(SallyDBOptions::RocksDB((b.0, b.1, RocksDBAccessType::Secondary(Some(q)), b.3, None))),
862                            None => {
863                                let p: std::path::PathBuf = tempfile::tempdir()
864                                    .expect("Failed to open temporary directory")
865                                    .into_path();
866                                #intermediate_db_map_struct_name::init(SallyDBOptions::RocksDB((b.0, b.1, RocksDBAccessType::Secondary(Some(p)), b.3, None)))
867                            }
868                        };
869                        Self {
870                            #(
871                                #field_names: inner.#field_names,
872                            )*
873                        }
874                    }
875                }
876            }
877
878            /// Dump all key-value pairs in the page at the given table name
879            /// Tables must be opened in read only mode using `open_tables_read_only`
880            pub fn dump(&self, table_name: &str, page_size: u16,
881                page_number: usize) -> eyre::Result<std::collections::BTreeMap<String, String>> {
882                Ok(match table_name {
883                    #(
884                        stringify!(#field_names) => {
885                            match &self.#field_names {
886                                SallyColumn::RocksDB((db_map, typed_store::sally::SallyConfig { mode: typed_store::sally::SallyRunMode::FallbackToDB })) => {
887                                    typed_store::traits::Map::try_catch_up_with_primary(db_map)?;
888                                    typed_store::traits::Map::unbounded_iter(db_map)
889                                        .skip((page_number * (page_size) as usize))
890                                        .take(page_size as usize)
891                                        .map(|(k, v)| (format!("{:?}", k), format!("{:?}", v)))
892                                        .collect::<std::collections::BTreeMap<_, _>>()
893                                }
894                                _ => unimplemented!(),
895                            }
896                        }
897                    )*
898                    _ => eyre::bail!("No such table name: {}", table_name),
899                })
900            }
901
902            pub fn table_summary(&self, table_name: &str) -> eyre::Result<typed_store::traits::TableSummary> {
903                let mut count = 0;
904                let mut key_bytes = 0;
905                let mut value_bytes = 0;
906                match table_name {
907                    #(
908                        stringify!(#field_names) => {
909                            match &self.#field_names {
910                                SallyColumn::RocksDB((db_map, typed_store::sally::SallyConfig { mode: typed_store::sally::SallyRunMode::FallbackToDB })) => {
911                                    typed_store::traits::Map::try_catch_up_with_primary(db_map)?;
912                                    db_map.table_summary()
913                                }
914                                _ => unimplemented!(),
915                            }
916                        }
917                    )*
918
919                    _ => eyre::bail!("No such table name: {}", table_name),
920                }
921            }
922
923            /// Count the keys in this table
924            /// Tables must be opened in read only mode using `open_tables_read_only`
925            pub fn count_keys(&self, table_name: &str) -> eyre::Result<usize> {
926                Ok(match table_name {
927                    #(
928                        stringify!(#field_names) => {
929                            match &self.#field_names {
930                                SallyColumn::RocksDB((db_map, typed_store::sally::SallyConfig { mode: typed_store::sally::SallyRunMode::FallbackToDB })) => {
931                                    typed_store::traits::Map::try_catch_up_with_primary(db_map)?;
932                                    typed_store::traits::Map::unbounded_iter(db_map).count()
933                                }
934                                _ => unimplemented!(),
935                            }
936                        }
937                    )*
938
939                    _ => eyre::bail!("No such table name: {}", table_name),
940                })
941            }
942
943            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
944                vec![#(
945                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
946                )*].into_iter().collect()
947            }
948        }
949
950
951        impl <
952                #(
953                    #generics_names: #generics_bounds_token,
954                )*
955            > TypedStoreDebug for #secondary_db_map_struct_name #generics {
956                fn dump_table(
957                    &self,
958                    table_name: String,
959                    page_size: u16,
960                    page_number: usize,
961                ) -> eyre::Result<std::collections::BTreeMap<String, String>> {
962                    self.dump(table_name.as_str(), page_size, page_number)
963                }
964
965                fn primary_db_name(&self) -> String {
966                    stringify!(#name).to_owned()
967                }
968
969                fn describe_all_tables(&self) -> std::collections::BTreeMap<String, (String, String)> {
970                    Self::describe_tables()
971                }
972
973                fn count_table_keys(&self, table_name: String) -> eyre::Result<usize> {
974                    self.count_keys(table_name.as_str())
975                }
976                fn table_summary(&self, table_name: String) -> eyre::Result<TableSummary> {
977                    self.table_summary(table_name.as_str())
978                }
979
980        }
981
982    })
983}
typed_store_derive/lib.rs

typed_store_derive/
lib.rs
