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 (db, rwopt_cfs) = {
327                    let opt_cfs = match tables_db_options_override {
328                        None => [
329                            #(
330                                (stringify!(#cf_names).to_owned(), #default_options_override_fn_names()),
331                            )*
332                        ],
333                        Some(o) => [
334                            #(
335                                (stringify!(#cf_names).to_owned(), o.to_map().get(stringify!(#cf_names)).unwrap().clone()),
336                            )*
337                        ]
338                    };
339                    // Safe to call unwrap because we will have at least one field_name entry in the struct
340                    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();
341                    let opt_cfs: Vec<_> = opt_cfs.iter().map(|q| (q.0.as_str(), q.1.options.clone())).collect();
342                    let db = match (as_secondary_with_path.clone(), is_transaction) {
343                        (Some(p), _) => typed_store::rocks::open_cf_opts_secondary(path, Some(&p), global_db_options_override, metric_conf, &opt_cfs),
344                        (_, true) => typed_store::rocks::open_cf_opts_transactional(path, global_db_options_override, metric_conf, &opt_cfs),
345                        _ => typed_store::rocks::open_cf_opts(path, global_db_options_override, metric_conf, &opt_cfs)
346                    };
347                    db.map(|d| (d, rwopt_cfs))
348                }.expect(&format!("Cannot open DB at {:?}", path));
349                let deprecated_tables = vec![#(stringify!(#deprecated_cfs),)*];
350                let (
351                        #(
352                            #field_names
353                        ),*
354                ) = (#(
355                        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))[..])
356                    ),*);
357
358                if as_secondary_with_path.is_none() && remove_deprecated_tables {
359                    #(
360                        db.drop_cf(stringify!(#deprecated_cfs)).expect("failed to drop a deprecated cf");
361                    )*
362                }
363                Self {
364                    #(
365                        #field_names,
366                    )*
367                }
368            }
369        }
370
371
372        // <----------- This section generates the read-write open logic and other common utils -------------->
373
374        impl <
375                #(
376                    #generics_names: #generics_bounds_token,
377                )*
378            > #name #generics {
379            /// Opens a set of tables in read-write mode
380            /// Only one process is allowed to do this at a time
381            /// `global_db_options_override` apply to the whole DB
382            /// `tables_db_options_override` apply to each table. If `None`, the attributes from `default_options_override_fn` are used if any
383            #[expect(unused_parens)]
384            pub fn open_tables_read_write(
385                path: std::path::PathBuf,
386                metric_conf: typed_store::rocks::MetricConf,
387                global_db_options_override: Option<typed_store::rocksdb::Options>,
388                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>
389            ) -> Self {
390                let inner = #intermediate_db_map_struct_name::open_tables_impl(path, None, false, metric_conf, global_db_options_override, tables_db_options_override, false);
391                Self {
392                    #(
393                        #field_names: #post_process_fn(inner.#field_names),
394                    )*
395                }
396            }
397
398            #[expect(unused_parens)]
399            pub fn open_tables_read_write_with_deprecation_option(
400                path: std::path::PathBuf,
401                metric_conf: typed_store::rocks::MetricConf,
402                global_db_options_override: Option<typed_store::rocksdb::Options>,
403                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>,
404                remove_deprecated_tables: bool,
405            ) -> Self {
406                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);
407                Self {
408                    #(
409                        #field_names: #post_process_fn(inner.#field_names),
410                    )*
411                }
412            }
413
414            /// Opens a set of tables in transactional read-write mode
415            /// Only one process is allowed to do this at a time
416            /// `global_db_options_override` apply to the whole DB
417            /// `tables_db_options_override` apply to each table. If `None`, the attributes from `default_options_override_fn` are used if any
418            #[expect(unused_parens)]
419            pub fn open_tables_transactional(
420                path: std::path::PathBuf,
421                metric_conf: typed_store::rocks::MetricConf,
422                global_db_options_override: Option<typed_store::rocksdb::Options>,
423                tables_db_options_override: Option<typed_store::rocks::DBMapTableConfigMap>
424            ) -> Self {
425                let inner = #intermediate_db_map_struct_name::open_tables_impl(path, None, true, metric_conf, global_db_options_override, tables_db_options_override, false);
426                Self {
427                    #(
428                        #field_names: #post_process_fn(inner.#field_names),
429                    )*
430                }
431            }
432
433            /// Returns a list of the tables name and type pairs
434            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
435                vec![#(
436                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
437                )*].into_iter().collect()
438            }
439
440            /// This opens the DB in read only mode and returns a struct which exposes debug features
441            pub fn get_read_only_handle (
442                primary_path: std::path::PathBuf,
443                with_secondary_path: Option<std::path::PathBuf>,
444                global_db_options_override: Option<typed_store::rocksdb::Options>,
445                metric_conf: typed_store::rocks::MetricConf,
446                ) -> #secondary_db_map_struct_name #generics {
447                #secondary_db_map_struct_name::open_tables_read_only(primary_path, with_secondary_path, metric_conf, global_db_options_override)
448            }
449        }
450
451
452        // <----------- This section generates the features that use read-only open logic -------------->
453        /// Create an intermediate struct used to open the DBMap tables in secondary mode
454        /// This is only used internally
455        pub struct #secondary_db_map_struct_name #generics {
456            #(
457                pub #field_names : DBMap #inner_types,
458            )*
459        }
460
461        impl <
462                #(
463                    #generics_names: #generics_bounds_token,
464                )*
465            > #secondary_db_map_struct_name #generics {
466            /// Open in read only mode. No limitation on number of processes to do this
467            pub fn open_tables_read_only(
468                primary_path: std::path::PathBuf,
469                with_secondary_path: Option<std::path::PathBuf>,
470                metric_conf: typed_store::rocks::MetricConf,
471                global_db_options_override: Option<typed_store::rocksdb::Options>,
472            ) -> Self {
473                let inner = match with_secondary_path {
474                    Some(q) => #intermediate_db_map_struct_name::open_tables_impl(primary_path, Some(q), false, metric_conf, global_db_options_override, None, false),
475                    None => {
476                        let p: std::path::PathBuf = tempfile::tempdir()
477                        .expect("Failed to open temporary directory")
478                        .into_path();
479                        #intermediate_db_map_struct_name::open_tables_impl(primary_path, Some(p), false, metric_conf, global_db_options_override, None, false)
480                    }
481                };
482                Self {
483                    #(
484                        #field_names: inner.#field_names,
485                    )*
486                }
487            }
488
489            fn cf_name_to_table_name(cf_name: &str) -> eyre::Result<&'static str> {
490                Ok(match cf_name {
491                    #(
492                        stringify!(#cf_names) => stringify!(#field_names),
493                    )*
494                    _ => eyre::bail!("No such cf name: {}", cf_name),
495                })
496            }
497
498            /// Dump all key-value pairs in the page at the given table name
499            /// Tables must be opened in read only mode using `open_tables_read_only`
500            pub fn dump(&self, cf_name: &str, page_size: u16, page_number: usize) -> eyre::Result<std::collections::BTreeMap<String, String>> {
501                let table_name = Self::cf_name_to_table_name(cf_name)?;
502
503                Ok(match table_name {
504                    #(
505                        stringify!(#field_names) => {
506                            typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
507                            typed_store::traits::Map::unbounded_iter(&self.#field_names)
508                                .skip((page_number * (page_size) as usize))
509                                .take(page_size as usize)
510                                .map(|(k, v)| (format!("{:?}", k), format!("{:?}", v)))
511                                .collect::<std::collections::BTreeMap<_, _>>()
512                        }
513                    )*
514
515                    _ => eyre::bail!("No such table name: {}", table_name),
516                })
517            }
518
519            /// Get key value sizes from the db
520            /// Tables must be opened in read only mode using `open_tables_read_only`
521            pub fn table_summary(&self, table_name: &str) -> eyre::Result<typed_store::traits::TableSummary> {
522                let mut count = 0;
523                let mut key_bytes = 0;
524                let mut value_bytes = 0;
525                match table_name {
526                    #(
527                        stringify!(#field_names) => {
528                            typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
529                            self.#field_names.table_summary()
530                        }
531                    )*
532
533                    _ => eyre::bail!("No such table name: {}", table_name),
534                }
535            }
536
537            /// Count the keys in this table
538            /// Tables must be opened in read only mode using `open_tables_read_only`
539            pub fn count_keys(&self, table_name: &str) -> eyre::Result<usize> {
540                Ok(match table_name {
541                    #(
542                        stringify!(#field_names) => {
543                            typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
544                            typed_store::traits::Map::unbounded_iter(&self.#field_names).count()
545                        }
546                    )*
547
548                    _ => eyre::bail!("No such table name: {}", table_name),
549                })
550            }
551
552            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
553                vec![#(
554                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
555                )*].into_iter().collect()
556            }
557
558            /// Try catch up with primary for all tables. This can be a slow operation
559            /// Tables must be opened in read only mode using `open_tables_read_only`
560            pub fn try_catch_up_with_primary_all(&self) -> eyre::Result<()> {
561                #(
562                    typed_store::traits::Map::try_catch_up_with_primary(&self.#field_names)?;
563                )*
564                Ok(())
565            }
566        }
567
568        impl <
569                #(
570                    #generics_names: #generics_bounds_token,
571                )*
572            > TypedStoreDebug for #secondary_db_map_struct_name #generics {
573                fn dump_table(
574                    &self,
575                    table_name: String,
576                    page_size: u16,
577                    page_number: usize,
578                ) -> eyre::Result<std::collections::BTreeMap<String, String>> {
579                    self.dump(table_name.as_str(), page_size, page_number)
580                }
581
582                fn primary_db_name(&self) -> String {
583                    stringify!(#name).to_owned()
584                }
585
586                fn describe_all_tables(&self) -> std::collections::BTreeMap<String, (String, String)> {
587                    Self::describe_tables()
588                }
589
590                fn count_table_keys(&self, table_name: String) -> eyre::Result<usize> {
591                    self.count_keys(table_name.as_str())
592                }
593
594                fn table_summary(&self, table_name: String) -> eyre::Result<TableSummary> {
595                    self.table_summary(table_name.as_str())
596                }
597
598
599        }
600
601    })
602}
603
604#[proc_macro_derive(SallyDB, attributes(default_options_override_fn))]
605pub fn derive_sallydb_general(input: TokenStream) -> TokenStream {
606    // log_syntax!("here");
607    let input = parse_macro_input!(input as ItemStruct);
608    let name = &input.ident;
609    let generics = &input.generics;
610    let generics_names = extract_generics_names(generics);
611
612    let allowed_types_with_post_process_fn: BTreeMap<_, _> =
613        [("SallyColumn", "")].into_iter().collect();
614    let allowed_strs = allowed_types_with_post_process_fn
615        .keys()
616        .map(|s| s.to_string())
617        .collect();
618
619    // TODO: use `parse_quote` over `parse()`
620    // TODO: Eventually this should return a Vec<Vec<GeneralTableOptions>> to
621    // capture default table options for each column type i.e. RockDB, TestDB, etc
622    let ExtractedStructInfo {
623        field_names,
624        inner_types,
625        derived_table_options,
626        simple_field_type_name_str,
627        ..
628    } = extract_struct_info(input.clone(), allowed_strs);
629
630    let (key_names, value_names): (Vec<_>, Vec<_>) = inner_types
631        .iter()
632        .map(|q| (q.args.first().unwrap(), q.args.last().unwrap()))
633        .unzip();
634
635    // This is the actual name of the type which was found
636    let post_process_fn_str = allowed_types_with_post_process_fn
637        .get(&simple_field_type_name_str.as_str())
638        .unwrap();
639    let post_process_fn: proc_macro2::TokenStream = post_process_fn_str.parse().unwrap();
640
641    let default_options_override_fn_names: Vec<proc_macro2::TokenStream> = derived_table_options
642        .iter()
643        .map(|q| {
644            let GeneralTableOptions::OverrideFunction(fn_name) = q;
645            fn_name.parse().unwrap()
646        })
647        .collect();
648
649    let generics_bounds =
650        "std::fmt::Debug + serde::Serialize + for<'de> serde::de::Deserialize<'de>";
651    let generics_bounds_token: proc_macro2::TokenStream = generics_bounds.parse().unwrap();
652
653    let config_struct_name_str = format!("{name}SallyConfigurator");
654    let sally_config_struct_name: proc_macro2::TokenStream =
655        config_struct_name_str.parse().unwrap();
656
657    let intermediate_db_map_struct_name_str = format!("{name}Primary");
658    let intermediate_db_map_struct_name: proc_macro2::TokenStream =
659        intermediate_db_map_struct_name_str.parse().unwrap();
660
661    let secondary_db_map_struct_name_str = format!("{name}ReadOnly");
662    let secondary_db_map_struct_name: proc_macro2::TokenStream =
663        secondary_db_map_struct_name_str.parse().unwrap();
664
665    TokenStream::from(quote! {
666
667        // <----------- This section generates the configurator struct -------------->
668
669        /// Create config structs for configuring SallyColumns
670        pub struct #sally_config_struct_name {
671            #(
672                pub #field_names : typed_store::sally::SallyColumnOptions,
673            )*
674        }
675
676        impl #sally_config_struct_name {
677            /// Initialize to defaults
678            pub fn init() -> Self {
679                Self {
680                    #(
681                        #field_names : typed_store::sally::default_column_options(),
682                    )*
683                }
684            }
685
686            /// Build a config
687            pub fn build(&self) -> typed_store::sally::SallyDBConfigMap {
688                typed_store::sally::SallyDBConfigMap::new([
689                    #(
690                        (stringify!(#field_names).to_owned(), self.#field_names.clone()),
691                    )*
692                ].into_iter().collect())
693            }
694        }
695
696
697        impl <
698                #(
699                    #generics_names: #generics_bounds_token,
700                )*
701            > #name #generics {
702
703                pub fn configurator() -> #sally_config_struct_name {
704                    #sally_config_struct_name::init()
705                }
706        }
707
708
709        // <----------- This section generates the core open logic for opening sally columns -------------->
710
711        /// Create an intermediate struct used to open the DBMap tables in primary mode
712        /// This is only used internally
713        struct #intermediate_db_map_struct_name #generics {
714                #(
715                    pub #field_names : SallyColumn #inner_types,
716                )*
717        }
718
719
720        impl <
721                #(
722                    #generics_names: #generics_bounds_token,
723                )*
724            > #intermediate_db_map_struct_name #generics {
725            /// Opens a set of tables in read-write mode
726            /// If as_secondary_with_path is set, the DB is opened in read only mode with the path specified
727            pub fn init(db_options: typed_store::sally::SallyDBOptions) -> Self {
728                match db_options {
729                    typed_store::sally::SallyDBOptions::TestDB => {
730                        let (
731                            #(
732                                #field_names
733                            ),*
734                        ) = (#(
735                            SallyColumn::TestDB((typed_store::test_db::TestDB::#inner_types::open(), typed_store::sally::SallyConfig::default()))
736                            ),*);
737
738                        Self {
739                            #(
740                                #field_names,
741                            )*
742                        }
743                    },
744                    typed_store::sally::SallyDBOptions::RocksDB((path, metric_conf, access_type, global_db_options_override, tables_db_options_override)) => {
745                        let path = &path;
746                        let (db, rwopt_cfs) = {
747                            let opt_cfs = match tables_db_options_override {
748                                None => [
749                                    #(
750                                        (stringify!(#field_names).to_owned(), #default_options_override_fn_names().clone()),
751                                    )*
752                                ],
753                                Some(o) => [
754                                    #(
755                                        (stringify!(#field_names).to_owned(), o.to_map().get(stringify!(#field_names)).unwrap().clone()),
756                                    )*
757                                ]
758                            };
759                            // Safe to call unwrap because we will have at least one field_name entry in the struct
760                            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();
761                            let opt_cfs: Vec<_> = opt_cfs.iter().map(|q| (q.0.as_str(), q.1.options.clone())).collect();
762                            let db = match access_type {
763                                RocksDBAccessType::Secondary(Some(p)) => typed_store::rocks::open_cf_opts_secondary(path, Some(&p), global_db_options_override, metric_conf, &opt_cfs),
764                                _ => typed_store::rocks::open_cf_opts(path, global_db_options_override, metric_conf, &opt_cfs)
765                            };
766                            db.map(|d| (d, rwopt_cfs))
767                        }.expect(&format!("Cannot open DB at {:?}", path));
768                        let (
769                            #(
770                                #field_names
771                            ),*
772                        ) = (#(
773                            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()))
774                            ),*);
775
776                        Self {
777                            #(
778                                #field_names,
779                            )*
780                        }
781                    }
782                }
783            }
784        }
785
786
787        // <----------- This section generates the read-write open logic and other common utils -------------->
788        impl <
789                #(
790                    #generics_names: #generics_bounds_token,
791                )*
792            > #name #generics {
793            /// Opens a set of tables in read-write mode
794            /// Only one process is allowed to do this at a time
795            /// `global_db_options_override` apply to the whole DB
796            /// `tables_db_options_override` apply to each table. If `None`, the attributes from `default_options_override_fn` are used if any
797            #[expect(unused_parens)]
798            pub fn init(
799                db_options: typed_store::sally::SallyDBOptions
800            ) -> Self {
801                let inner = #intermediate_db_map_struct_name::init(db_options);
802                Self {
803                    #(
804                        #field_names: #post_process_fn(inner.#field_names),
805                    )*
806                }
807            }
808
809            /// Returns a list of the tables name and type pairs
810            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
811                vec![#(
812                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
813                )*].into_iter().collect()
814            }
815
816            /// This opens the DB in read only mode and returns a struct which exposes debug features
817            pub fn get_read_only_handle (
818                db_options: typed_store::sally::SallyReadOnlyDBOptions
819                ) -> #secondary_db_map_struct_name #generics {
820                #secondary_db_map_struct_name::init_read_only(db_options)
821            }
822        }
823
824        // <----------- This section generates the features that use read-only open logic -------------->
825        /// Create an intermediate struct used to open the DBMap tables in secondary mode
826        /// This is only used internally
827        pub struct #secondary_db_map_struct_name #generics {
828            #(
829                pub #field_names : SallyColumn #inner_types,
830            )*
831        }
832
833        impl <
834                #(
835                    #generics_names: #generics_bounds_token,
836                )*
837            > #secondary_db_map_struct_name #generics {
838            /// Open in read only mode. No limitation on number of processes to do this
839            pub fn init_read_only(
840                db_options: typed_store::sally::SallyReadOnlyDBOptions,
841            ) -> Self {
842                match db_options {
843                    typed_store::sally::SallyReadOnlyDBOptions::TestDB => {
844                        let inner = #intermediate_db_map_struct_name::init(SallyDBOptions::TestDB);
845                        Self {
846                            #(
847                                #field_names: inner.#field_names,
848                            )*
849                        }
850                    },
851                    typed_store::sally::SallyReadOnlyDBOptions::RocksDB(b) => {
852                        let inner = match b.2 {
853                            Some(q) => #intermediate_db_map_struct_name::init(SallyDBOptions::RocksDB((b.0, b.1, RocksDBAccessType::Secondary(Some(q)), b.3, None))),
854                            None => {
855                                let p: std::path::PathBuf = tempfile::tempdir()
856                                    .expect("Failed to open temporary directory")
857                                    .into_path();
858                                #intermediate_db_map_struct_name::init(SallyDBOptions::RocksDB((b.0, b.1, RocksDBAccessType::Secondary(Some(p)), b.3, None)))
859                            }
860                        };
861                        Self {
862                            #(
863                                #field_names: inner.#field_names,
864                            )*
865                        }
866                    }
867                }
868            }
869
870            /// Dump all key-value pairs in the page at the given table name
871            /// Tables must be opened in read only mode using `open_tables_read_only`
872            pub fn dump(&self, table_name: &str, page_size: u16,
873                page_number: usize) -> eyre::Result<std::collections::BTreeMap<String, String>> {
874                Ok(match table_name {
875                    #(
876                        stringify!(#field_names) => {
877                            match &self.#field_names {
878                                SallyColumn::RocksDB((db_map, typed_store::sally::SallyConfig { mode: typed_store::sally::SallyRunMode::FallbackToDB })) => {
879                                    typed_store::traits::Map::try_catch_up_with_primary(db_map)?;
880                                    typed_store::traits::Map::unbounded_iter(db_map)
881                                        .skip((page_number * (page_size) as usize))
882                                        .take(page_size as usize)
883                                        .map(|(k, v)| (format!("{:?}", k), format!("{:?}", v)))
884                                        .collect::<std::collections::BTreeMap<_, _>>()
885                                }
886                                _ => unimplemented!(),
887                            }
888                        }
889                    )*
890                    _ => eyre::bail!("No such table name: {}", table_name),
891                })
892            }
893
894            pub fn table_summary(&self, table_name: &str) -> eyre::Result<typed_store::traits::TableSummary> {
895                let mut count = 0;
896                let mut key_bytes = 0;
897                let mut value_bytes = 0;
898                match table_name {
899                    #(
900                        stringify!(#field_names) => {
901                            match &self.#field_names {
902                                SallyColumn::RocksDB((db_map, typed_store::sally::SallyConfig { mode: typed_store::sally::SallyRunMode::FallbackToDB })) => {
903                                    typed_store::traits::Map::try_catch_up_with_primary(db_map)?;
904                                    db_map.table_summary()
905                                }
906                                _ => unimplemented!(),
907                            }
908                        }
909                    )*
910
911                    _ => eyre::bail!("No such table name: {}", table_name),
912                }
913            }
914
915            /// Count the keys in this table
916            /// Tables must be opened in read only mode using `open_tables_read_only`
917            pub fn count_keys(&self, table_name: &str) -> eyre::Result<usize> {
918                Ok(match table_name {
919                    #(
920                        stringify!(#field_names) => {
921                            match &self.#field_names {
922                                SallyColumn::RocksDB((db_map, typed_store::sally::SallyConfig { mode: typed_store::sally::SallyRunMode::FallbackToDB })) => {
923                                    typed_store::traits::Map::try_catch_up_with_primary(db_map)?;
924                                    typed_store::traits::Map::unbounded_iter(db_map).count()
925                                }
926                                _ => unimplemented!(),
927                            }
928                        }
929                    )*
930
931                    _ => eyre::bail!("No such table name: {}", table_name),
932                })
933            }
934
935            pub fn describe_tables() -> std::collections::BTreeMap<String, (String, String)> {
936                vec![#(
937                    (stringify!(#field_names).to_owned(), (stringify!(#key_names).to_owned(), stringify!(#value_names).to_owned())),
938                )*].into_iter().collect()
939            }
940        }
941
942
943        impl <
944                #(
945                    #generics_names: #generics_bounds_token,
946                )*
947            > TypedStoreDebug for #secondary_db_map_struct_name #generics {
948                fn dump_table(
949                    &self,
950                    table_name: String,
951                    page_size: u16,
952                    page_number: usize,
953                ) -> eyre::Result<std::collections::BTreeMap<String, String>> {
954                    self.dump(table_name.as_str(), page_size, page_number)
955                }
956
957                fn primary_db_name(&self) -> String {
958                    stringify!(#name).to_owned()
959                }
960
961                fn describe_all_tables(&self) -> std::collections::BTreeMap<String, (String, String)> {
962                    Self::describe_tables()
963                }
964
965                fn count_table_keys(&self, table_name: String) -> eyre::Result<usize> {
966                    self.count_keys(table_name.as_str())
967                }
968                fn table_summary(&self, table_name: String) -> eyre::Result<TableSummary> {
969                    self.table_summary(table_name.as_str())
970                }
971
972        }
973
974    })
975}
typed_store_derive/lib.rs

typed_store_derive/
lib.rs
