1mod crate_path;
7mod db;
8
9use std::sync::Arc;
10
11use crate::{
12 build::get_schema,
13 node::{Canister, Entity, Schema, Store},
14};
15use icydb_schema::encode_schema_fragment;
16use proc_macro2::TokenStream;
17use quote::quote;
18use sha2::{Digest, Sha256};
19
20#[must_use]
28pub fn generate_with_options(canister_path: &str, options: BuildOptions) -> String {
29 let schema = get_schema().expect("schema must be valid before codegen");
31 let canister = schema
32 .cast_node::<Canister>(canister_path)
33 .expect("canister path must resolve to a canister node");
34 let fragment = schema
35 .schema_fragment_for_canister(canister_path)
36 .expect("sealed canister database closure must lower into a schema fragment");
37
38 let code = ActorBuilder::new(
40 Arc::new(schema.clone()),
41 canister.clone(),
42 options,
43 fragment,
44 );
45 drop(schema);
46 let tokens = crate_path::rewrite_icydb_path(code.generate(), options.icydb_crate_path());
47
48 tokens.to_string()
49}
50
51#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
59pub struct BuildOptions {
60 sql: BuildSqlOptions,
61 metrics: BuildMetricsOptions,
62 snapshot_enabled: bool,
63 schema_enabled: bool,
64 icydb_crate_path: Option<&'static str>,
65}
66
67impl BuildOptions {
68 #[must_use]
70 pub const fn with_sql_readonly_enabled(mut self, enabled: bool) -> Self {
71 self.sql.surfaces = self.sql.surfaces.with_readonly_enabled(enabled);
72
73 self
74 }
75
76 #[must_use]
78 pub const fn with_sql_ddl_enabled(mut self, enabled: bool) -> Self {
79 self.sql.surfaces = self.sql.surfaces.with_ddl_enabled(enabled);
80
81 self
82 }
83
84 #[must_use]
86 pub const fn with_sql_fixtures_enabled(mut self, enabled: bool) -> Self {
87 self.sql.surfaces = self.sql.surfaces.with_fixtures_enabled(enabled);
88
89 self
90 }
91
92 #[must_use]
94 pub const fn with_sql_integrity_enabled(mut self, enabled: bool) -> Self {
95 self.sql.surfaces = self.sql.surfaces.with_integrity_enabled(enabled);
96
97 self
98 }
99
100 #[must_use]
102 pub const fn with_sql_introspection_enabled(mut self, enabled: bool) -> Self {
103 self.sql.surfaces = self.sql.surfaces.with_introspection_enabled(enabled);
104
105 self
106 }
107
108 #[must_use]
110 pub const fn with_sql_update_policy(mut self, policy: Option<BuildSqlUpdatePolicy>) -> Self {
111 self.sql.update_policy = policy;
112
113 self
114 }
115
116 #[must_use]
118 pub const fn with_metrics_enabled(mut self, enabled: bool) -> Self {
119 self.metrics.enabled = enabled;
120
121 self
122 }
123
124 #[must_use]
126 pub const fn with_metrics_extended_enabled(mut self, enabled: bool) -> Self {
127 self.metrics.extended_enabled = enabled;
128
129 self
130 }
131
132 #[must_use]
134 pub const fn with_snapshot_enabled(mut self, enabled: bool) -> Self {
135 self.snapshot_enabled = enabled;
136
137 self
138 }
139
140 #[must_use]
142 pub const fn with_schema_enabled(mut self, enabled: bool) -> Self {
143 self.schema_enabled = enabled;
144
145 self
146 }
147
148 #[must_use]
154 pub const fn with_icydb_crate_path(mut self, path: &'static str) -> Self {
155 self.icydb_crate_path = Some(path);
156
157 self
158 }
159
160 #[must_use]
162 pub const fn sql_readonly_enabled(self) -> bool {
163 self.sql.surfaces.readonly_enabled()
164 }
165
166 #[must_use]
168 pub const fn sql_ddl_enabled(self) -> bool {
169 self.sql.surfaces.ddl_enabled()
170 }
171
172 #[must_use]
174 pub const fn sql_fixtures_enabled(self) -> bool {
175 self.sql.surfaces.fixtures_enabled()
176 }
177
178 #[must_use]
180 pub const fn sql_integrity_enabled(self) -> bool {
181 self.sql.surfaces.integrity_enabled()
182 }
183
184 #[must_use]
186 pub const fn sql_introspection_enabled(self) -> bool {
187 self.sql.surfaces.introspection_enabled()
188 }
189
190 #[must_use]
191 pub(crate) const fn sql_surface_flags(self) -> BuildSqlSurfaceFlags {
192 self.sql.surfaces
193 }
194
195 #[must_use]
196 const fn icydb_crate_path(self) -> Option<&'static str> {
197 self.icydb_crate_path
198 }
199
200 #[must_use]
202 pub const fn sql_update_policy(self) -> Option<BuildSqlUpdatePolicy> {
203 self.sql.update_policy
204 }
205
206 #[must_use]
208 pub const fn sql_update_enabled(self) -> bool {
209 self.sql_update_policy().is_some()
210 }
211
212 #[must_use]
214 pub const fn metrics_enabled(self) -> bool {
215 self.metrics.enabled
216 }
217
218 #[must_use]
220 pub const fn metrics_extended_enabled(self) -> bool {
221 self.metrics.enabled && self.metrics.extended_enabled
222 }
223
224 #[must_use]
226 pub const fn snapshot_enabled(self) -> bool {
227 self.snapshot_enabled
228 }
229
230 #[must_use]
232 pub const fn schema_enabled(self) -> bool {
233 self.schema_enabled
234 }
235
236 #[must_use]
238 pub const fn sql_enabled(self) -> bool {
239 self.sql_readonly_enabled()
240 || self.sql_ddl_enabled()
241 || self.sql_fixtures_enabled()
242 || self.sql_integrity_enabled()
243 || self.sql_update_enabled()
244 }
245}
246
247#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
248struct BuildSqlOptions {
249 surfaces: BuildSqlSurfaceFlags,
250 update_policy: Option<BuildSqlUpdatePolicy>,
251}
252
253#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
254pub(crate) struct BuildSqlSurfaceFlags(u8);
255
256impl BuildSqlSurfaceFlags {
257 const DDL: u8 = 1 << 1;
258 const FIXTURES: u8 = 1 << 2;
259 const INTROSPECTION: u8 = 1 << 3;
260 const INTEGRITY: u8 = 1 << 4;
261 const READONLY: u8 = 1;
262
263 #[must_use]
264 pub(crate) const fn with_readonly_enabled(self, enabled: bool) -> Self {
265 self.with_flag(Self::READONLY, enabled)
266 }
267
268 #[must_use]
269 pub(crate) const fn with_ddl_enabled(self, enabled: bool) -> Self {
270 self.with_flag(Self::DDL, enabled)
271 }
272
273 #[must_use]
274 pub(crate) const fn with_fixtures_enabled(self, enabled: bool) -> Self {
275 self.with_flag(Self::FIXTURES, enabled)
276 }
277
278 #[must_use]
279 pub(crate) const fn with_integrity_enabled(self, enabled: bool) -> Self {
280 self.with_flag(Self::INTEGRITY, enabled)
281 }
282
283 #[must_use]
284 pub(crate) const fn with_introspection_enabled(self, enabled: bool) -> Self {
285 self.with_flag(Self::INTROSPECTION, enabled)
286 }
287
288 #[must_use]
289 pub(crate) const fn readonly_enabled(self) -> bool {
290 self.contains(Self::READONLY)
291 }
292
293 #[must_use]
294 pub(crate) const fn ddl_enabled(self) -> bool {
295 self.contains(Self::DDL)
296 }
297
298 #[must_use]
299 pub(crate) const fn fixtures_enabled(self) -> bool {
300 self.contains(Self::FIXTURES)
301 }
302
303 #[must_use]
304 pub(crate) const fn integrity_enabled(self) -> bool {
305 self.contains(Self::INTEGRITY)
306 }
307
308 #[must_use]
309 pub(crate) const fn introspection_enabled(self) -> bool {
310 self.contains(Self::INTROSPECTION)
311 }
312
313 #[must_use]
314 const fn contains(self, flag: u8) -> bool {
315 self.0 & flag == flag
316 }
317
318 #[must_use]
319 const fn with_flag(self, flag: u8, enabled: bool) -> Self {
320 if enabled {
321 Self(self.0 | flag)
322 } else {
323 Self(self.0 & !flag)
324 }
325 }
326}
327
328#[derive(Clone, Copy, Debug, Eq, PartialEq)]
329struct BuildMetricsOptions {
330 enabled: bool,
331 extended_enabled: bool,
332}
333
334impl Default for BuildMetricsOptions {
335 fn default() -> Self {
336 Self {
337 enabled: true,
338 extended_enabled: false,
339 }
340 }
341}
342
343#[derive(Clone, Copy, Debug, Eq, PartialEq)]
346pub enum BuildSqlUpdatePolicy {
347 PublicPrimaryKeyOnly,
349 PublicBoundedDeterministic,
351}
352
353pub(crate) struct ActorBuilder {
361 pub(crate) schema: Arc<Schema>,
362 pub(crate) canister: Canister,
363 pub(crate) options: BuildOptions,
364 pub(crate) schema_fragment_bytes: Vec<u8>,
365 pub(crate) schema_submission_key: String,
366}
367
368impl ActorBuilder {
369 #[must_use]
371 pub fn new(
372 schema: Arc<Schema>,
373 canister: Canister,
374 options: BuildOptions,
375 fragment: icydb_schema::SchemaFragment,
376 ) -> Self {
377 let schema_fragment_bytes =
378 encode_schema_fragment(&fragment).expect("sealed schema fragment must encode");
379 let digest = Sha256::digest(schema_fragment_bytes.as_slice());
380 let schema_submission_key = format!("generated/{}", hex_bytes(digest.as_slice()));
381
382 Self {
383 schema,
384 canister,
385 options,
386 schema_fragment_bytes,
387 schema_submission_key,
388 }
389 }
390
391 #[must_use]
393 pub fn generate(self) -> TokenStream {
394 let mut tokens = quote!();
395
396 tokens.extend(db::generate(&self));
398 tokens.extend(generate_snapshot(&self));
399 tokens.extend(generate_metrics(&self));
400
401 quote! {
402 #tokens
403 }
404 }
405
406 #[must_use]
408 pub fn get_stores(&self) -> Vec<(String, Store)> {
409 let canister_path = self.canister_path();
410
411 self.schema
412 .filter_nodes::<Store>(|node| node.canister() == canister_path)
413 .map(|(path, store)| (path.to_string(), store.clone()))
414 .collect()
415 }
416
417 #[must_use]
419 pub fn get_entities(&self) -> Vec<(String, Entity)> {
420 let canister_path = self.canister_path();
421
422 self.schema
423 .get_nodes::<Entity>()
424 .filter_map(|(path, entity)| {
425 let store = self.schema.cast_node::<Store>(entity.store()).ok()?;
426 if store.canister() == canister_path {
427 Some((path.to_string(), entity.clone()))
428 } else {
429 None
430 }
431 })
432 .collect()
433 }
434
435 fn canister_path(&self) -> String {
436 self.canister.def().path()
437 }
438}
439
440fn hex_bytes(bytes: &[u8]) -> String {
441 const HEX: &[u8; 16] = b"0123456789abcdef";
442 let mut encoded = String::with_capacity(bytes.len() * 2);
443 for byte in bytes {
444 encoded.push(char::from(HEX[usize::from(byte >> 4)]));
445 encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
446 }
447 encoded
448}
449
450#[must_use]
452fn generate_snapshot(builder: &ActorBuilder) -> TokenStream {
453 if builder.options.snapshot_enabled() {
454 quote! {
455 #[::icydb::__reexports::ic_cdk::query(name = "icydb_snapshot")]
456 pub fn __icydb_snapshot() -> Result<::icydb::db::StorageReport, ::icydb::Error> {
457 ::icydb::__macro::execute_generated_storage_report(&db()?)
458 }
459 }
460 } else {
461 TokenStream::new()
462 }
463}
464
465#[must_use]
467fn generate_metrics(builder: &ActorBuilder) -> TokenStream {
468 let metrics_endpoint = builder.options.metrics_enabled().then(|| {
469 quote! {
470 #[::icydb::__reexports::ic_cdk::query(name = "icydb_metrics")]
471 pub fn __icydb_metrics(window_start_ms: Option<u64>) -> Result<::icydb::metrics::CompactMetricsReport, ::icydb::Error> {
472 Ok(::icydb::metrics::compact_metrics_report(window_start_ms))
473 }
474 }
475 });
476
477 let metrics_extended_endpoint = builder.options.metrics_extended_enabled().then(|| {
478 quote! {
479 #[::icydb::__reexports::ic_cdk::query(name = "icydb_metrics_extended")]
480 pub fn __icydb_metrics_extended(window_start_ms: Option<u64>) -> Result<::icydb::metrics::EventReport, ::icydb::Error> {
481 Ok(::icydb::metrics::metrics_report(window_start_ms))
482 }
483 }
484 });
485
486 let metrics_reset_endpoint = builder.options.metrics_enabled().then(|| {
487 quote! {
488 #[::icydb::__reexports::ic_cdk::update(name = "icydb_metrics_reset")]
489 pub fn __icydb_metrics_reset() -> Result<(), ::icydb::Error> {
490 ::icydb::metrics::metrics_reset_all();
491
492 Ok(())
493 }
494 }
495 });
496
497 quote! {
498 #metrics_endpoint
499 #metrics_extended_endpoint
500 #metrics_reset_endpoint
501 }
502}
503
504#[cfg(test)]
505mod tests {
506 use std::sync::Arc;
507
508 use crate::node::{Canister, Def, Schema};
509 use proc_macro2::TokenStream;
510
511 use super::{ActorBuilder, BuildOptions};
512
513 fn compact_tokens(tokens: TokenStream) -> String {
514 tokens
515 .to_string()
516 .chars()
517 .filter(|character| !character.is_whitespace())
518 .collect()
519 }
520
521 fn actor_builder_with_options(options: BuildOptions) -> ActorBuilder {
522 ActorBuilder::new(
523 Arc::new(Schema::new()),
524 Canister::new(Def::new("test", "Canister"), "test", 0, 1, 2, 3),
525 options,
526 icydb_schema::SchemaFragment::try_new(Vec::new(), Vec::new())
527 .expect("empty test fragment should admit"),
528 )
529 }
530
531 #[test]
532 fn default_build_options_enable_minimal_metrics_only() {
533 let options = BuildOptions::default();
534
535 assert!(!options.sql_readonly_enabled());
536 assert!(!options.sql_ddl_enabled());
537 assert!(!options.sql_fixtures_enabled());
538 assert!(!options.sql_integrity_enabled());
539 assert!(!options.sql_introspection_enabled());
540 assert!(!options.sql_update_enabled());
541 assert_eq!(options.sql_update_policy(), None);
542 assert!(options.metrics_enabled());
543 assert!(!options.metrics_extended_enabled());
544 assert!(!options.snapshot_enabled());
545 assert!(!options.schema_enabled());
546 }
547
548 #[test]
549 fn extended_metrics_requires_metrics_surface() {
550 let options = BuildOptions::default()
551 .with_metrics_enabled(false)
552 .with_metrics_extended_enabled(true);
553
554 assert!(!options.metrics_enabled());
555 assert!(!options.metrics_extended_enabled());
556
557 let options = options.with_metrics_enabled(true);
558
559 assert!(options.metrics_enabled());
560 assert!(options.metrics_extended_enabled());
561 }
562
563 #[test]
564 fn generated_metrics_surface_uses_public_icydb_endpoint_names() {
565 let builder =
566 actor_builder_with_options(BuildOptions::default().with_metrics_extended_enabled(true));
567 let surface = compact_tokens(super::generate_metrics(&builder));
568
569 assert!(surface.contains("name=\"icydb_metrics\""));
570 assert!(surface.contains("name=\"icydb_metrics_extended\""));
571 assert!(surface.contains("name=\"icydb_metrics_reset\""));
572 assert!(surface.contains("pubfn__icydb_metrics("));
573 assert!(surface.contains("pubfn__icydb_metrics_extended("));
574 assert!(surface.contains("pubfn__icydb_metrics_reset("));
575 }
576
577 #[test]
578 fn generated_snapshot_surface_uses_public_icydb_endpoint_name() {
579 let builder =
580 actor_builder_with_options(BuildOptions::default().with_snapshot_enabled(true));
581 let surface = compact_tokens(super::generate_snapshot(&builder));
582
583 assert!(surface.contains("name=\"icydb_snapshot\""));
584 assert!(surface.contains("pubfn__icydb_snapshot("));
585 }
586}