mod crate_path;
mod db;
use std::sync::Arc;
use crate::{
build::get_schema,
node::{Canister, Entity, Schema, Store},
};
use icydb_schema::encode_schema_fragment;
use proc_macro2::TokenStream;
use quote::quote;
use sha2::{Digest, Sha256};
#[must_use]
pub fn generate_with_options(canister_path: &str, options: BuildOptions) -> String {
let schema = get_schema().expect("schema must be valid before codegen");
let canister = schema
.cast_node::<Canister>(canister_path)
.expect("canister path must resolve to a canister node");
let fragment = schema
.schema_fragment_for_canister(canister_path)
.expect("sealed canister database closure must lower into a schema fragment");
let code = ActorBuilder::new(
Arc::new(schema.clone()),
canister.clone(),
options,
fragment,
);
drop(schema);
let tokens = crate_path::rewrite_icydb_path(code.generate(), options.icydb_crate_path());
tokens.to_string()
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct BuildOptions {
sql: BuildSqlOptions,
metrics: BuildMetricsOptions,
snapshot_enabled: bool,
schema_enabled: bool,
icydb_crate_path: Option<&'static str>,
}
impl BuildOptions {
#[must_use]
pub const fn with_sql_readonly_enabled(mut self, enabled: bool) -> Self {
self.sql.surfaces = self.sql.surfaces.with_readonly_enabled(enabled);
self
}
#[must_use]
pub const fn with_sql_ddl_enabled(mut self, enabled: bool) -> Self {
self.sql.surfaces = self.sql.surfaces.with_ddl_enabled(enabled);
self
}
#[must_use]
pub const fn with_sql_fixtures_enabled(mut self, enabled: bool) -> Self {
self.sql.surfaces = self.sql.surfaces.with_fixtures_enabled(enabled);
self
}
#[must_use]
pub const fn with_sql_integrity_enabled(mut self, enabled: bool) -> Self {
self.sql.surfaces = self.sql.surfaces.with_integrity_enabled(enabled);
self
}
#[must_use]
pub const fn with_sql_introspection_enabled(mut self, enabled: bool) -> Self {
self.sql.surfaces = self.sql.surfaces.with_introspection_enabled(enabled);
self
}
#[must_use]
pub const fn with_sql_update_policy(mut self, policy: Option<BuildSqlUpdatePolicy>) -> Self {
self.sql.update_policy = policy;
self
}
#[must_use]
pub const fn with_metrics_enabled(mut self, enabled: bool) -> Self {
self.metrics.enabled = enabled;
self
}
#[must_use]
pub const fn with_metrics_extended_enabled(mut self, enabled: bool) -> Self {
self.metrics.extended_enabled = enabled;
self
}
#[must_use]
pub const fn with_snapshot_enabled(mut self, enabled: bool) -> Self {
self.snapshot_enabled = enabled;
self
}
#[must_use]
pub const fn with_schema_enabled(mut self, enabled: bool) -> Self {
self.schema_enabled = enabled;
self
}
#[must_use]
pub const fn with_icydb_crate_path(mut self, path: &'static str) -> Self {
self.icydb_crate_path = Some(path);
self
}
#[must_use]
pub const fn sql_readonly_enabled(self) -> bool {
self.sql.surfaces.readonly_enabled()
}
#[must_use]
pub const fn sql_ddl_enabled(self) -> bool {
self.sql.surfaces.ddl_enabled()
}
#[must_use]
pub const fn sql_fixtures_enabled(self) -> bool {
self.sql.surfaces.fixtures_enabled()
}
#[must_use]
pub const fn sql_integrity_enabled(self) -> bool {
self.sql.surfaces.integrity_enabled()
}
#[must_use]
pub const fn sql_introspection_enabled(self) -> bool {
self.sql.surfaces.introspection_enabled()
}
#[must_use]
pub(crate) const fn sql_surface_flags(self) -> BuildSqlSurfaceFlags {
self.sql.surfaces
}
#[must_use]
const fn icydb_crate_path(self) -> Option<&'static str> {
self.icydb_crate_path
}
#[must_use]
pub const fn sql_update_policy(self) -> Option<BuildSqlUpdatePolicy> {
self.sql.update_policy
}
#[must_use]
pub const fn sql_update_enabled(self) -> bool {
self.sql_update_policy().is_some()
}
#[must_use]
pub const fn metrics_enabled(self) -> bool {
self.metrics.enabled
}
#[must_use]
pub const fn metrics_extended_enabled(self) -> bool {
self.metrics.enabled && self.metrics.extended_enabled
}
#[must_use]
pub const fn snapshot_enabled(self) -> bool {
self.snapshot_enabled
}
#[must_use]
pub const fn schema_enabled(self) -> bool {
self.schema_enabled
}
#[must_use]
pub const fn sql_enabled(self) -> bool {
self.sql_readonly_enabled()
|| self.sql_ddl_enabled()
|| self.sql_fixtures_enabled()
|| self.sql_integrity_enabled()
|| self.sql_update_enabled()
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
struct BuildSqlOptions {
surfaces: BuildSqlSurfaceFlags,
update_policy: Option<BuildSqlUpdatePolicy>,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(crate) struct BuildSqlSurfaceFlags(u8);
impl BuildSqlSurfaceFlags {
const DDL: u8 = 1 << 1;
const FIXTURES: u8 = 1 << 2;
const INTROSPECTION: u8 = 1 << 3;
const INTEGRITY: u8 = 1 << 4;
const READONLY: u8 = 1;
#[must_use]
pub(crate) const fn with_readonly_enabled(self, enabled: bool) -> Self {
self.with_flag(Self::READONLY, enabled)
}
#[must_use]
pub(crate) const fn with_ddl_enabled(self, enabled: bool) -> Self {
self.with_flag(Self::DDL, enabled)
}
#[must_use]
pub(crate) const fn with_fixtures_enabled(self, enabled: bool) -> Self {
self.with_flag(Self::FIXTURES, enabled)
}
#[must_use]
pub(crate) const fn with_integrity_enabled(self, enabled: bool) -> Self {
self.with_flag(Self::INTEGRITY, enabled)
}
#[must_use]
pub(crate) const fn with_introspection_enabled(self, enabled: bool) -> Self {
self.with_flag(Self::INTROSPECTION, enabled)
}
#[must_use]
pub(crate) const fn readonly_enabled(self) -> bool {
self.contains(Self::READONLY)
}
#[must_use]
pub(crate) const fn ddl_enabled(self) -> bool {
self.contains(Self::DDL)
}
#[must_use]
pub(crate) const fn fixtures_enabled(self) -> bool {
self.contains(Self::FIXTURES)
}
#[must_use]
pub(crate) const fn integrity_enabled(self) -> bool {
self.contains(Self::INTEGRITY)
}
#[must_use]
pub(crate) const fn introspection_enabled(self) -> bool {
self.contains(Self::INTROSPECTION)
}
#[must_use]
const fn contains(self, flag: u8) -> bool {
self.0 & flag == flag
}
#[must_use]
const fn with_flag(self, flag: u8, enabled: bool) -> Self {
if enabled {
Self(self.0 | flag)
} else {
Self(self.0 & !flag)
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct BuildMetricsOptions {
enabled: bool,
extended_enabled: bool,
}
impl Default for BuildMetricsOptions {
fn default() -> Self {
Self {
enabled: true,
extended_enabled: false,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BuildSqlUpdatePolicy {
PublicPrimaryKeyOnly,
PublicBoundedDeterministic,
}
pub(crate) struct ActorBuilder {
pub(crate) schema: Arc<Schema>,
pub(crate) canister: Canister,
pub(crate) options: BuildOptions,
pub(crate) schema_fragment_bytes: Vec<u8>,
pub(crate) schema_submission_key: String,
}
impl ActorBuilder {
#[must_use]
pub fn new(
schema: Arc<Schema>,
canister: Canister,
options: BuildOptions,
fragment: icydb_schema::SchemaFragment,
) -> Self {
let schema_fragment_bytes =
encode_schema_fragment(&fragment).expect("sealed schema fragment must encode");
let digest = Sha256::digest(schema_fragment_bytes.as_slice());
let schema_submission_key = format!("generated/{}", hex_bytes(digest.as_slice()));
Self {
schema,
canister,
options,
schema_fragment_bytes,
schema_submission_key,
}
}
#[must_use]
pub fn generate(self) -> TokenStream {
let mut tokens = quote!();
tokens.extend(db::generate(&self));
tokens.extend(generate_snapshot(&self));
tokens.extend(generate_metrics(&self));
quote! {
#tokens
}
}
#[must_use]
pub fn get_stores(&self) -> Vec<(String, Store)> {
let canister_path = self.canister_path();
self.schema
.filter_nodes::<Store>(|node| node.canister() == canister_path)
.map(|(path, store)| (path.to_string(), store.clone()))
.collect()
}
#[must_use]
pub fn get_entities(&self) -> Vec<(String, Entity)> {
let canister_path = self.canister_path();
self.schema
.get_nodes::<Entity>()
.filter_map(|(path, entity)| {
let store = self.schema.cast_node::<Store>(entity.store()).ok()?;
if store.canister() == canister_path {
Some((path.to_string(), entity.clone()))
} else {
None
}
})
.collect()
}
fn canister_path(&self) -> String {
self.canister.def().path()
}
}
fn hex_bytes(bytes: &[u8]) -> String {
const HEX: &[u8; 16] = b"0123456789abcdef";
let mut encoded = String::with_capacity(bytes.len() * 2);
for byte in bytes {
encoded.push(char::from(HEX[usize::from(byte >> 4)]));
encoded.push(char::from(HEX[usize::from(byte & 0x0f)]));
}
encoded
}
#[must_use]
fn generate_snapshot(builder: &ActorBuilder) -> TokenStream {
if builder.options.snapshot_enabled() {
quote! {
#[::icydb::__reexports::ic_cdk::query(name = "icydb_snapshot")]
pub fn __icydb_snapshot() -> Result<::icydb::db::StorageReport, ::icydb::Error> {
::icydb::__macro::execute_generated_storage_report(&db()?)
}
}
} else {
TokenStream::new()
}
}
#[must_use]
fn generate_metrics(builder: &ActorBuilder) -> TokenStream {
let metrics_endpoint = builder.options.metrics_enabled().then(|| {
quote! {
#[::icydb::__reexports::ic_cdk::query(name = "icydb_metrics")]
pub fn __icydb_metrics(window_start_ms: Option<u64>) -> Result<::icydb::metrics::CompactMetricsReport, ::icydb::Error> {
Ok(::icydb::metrics::compact_metrics_report(window_start_ms))
}
}
});
let metrics_extended_endpoint = builder.options.metrics_extended_enabled().then(|| {
quote! {
#[::icydb::__reexports::ic_cdk::query(name = "icydb_metrics_extended")]
pub fn __icydb_metrics_extended(window_start_ms: Option<u64>) -> Result<::icydb::metrics::EventReport, ::icydb::Error> {
Ok(::icydb::metrics::metrics_report(window_start_ms))
}
}
});
let metrics_reset_endpoint = builder.options.metrics_enabled().then(|| {
quote! {
#[::icydb::__reexports::ic_cdk::update(name = "icydb_metrics_reset")]
pub fn __icydb_metrics_reset() -> Result<(), ::icydb::Error> {
::icydb::metrics::metrics_reset_all();
Ok(())
}
}
});
quote! {
#metrics_endpoint
#metrics_extended_endpoint
#metrics_reset_endpoint
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use crate::node::{Canister, Def, Schema};
use proc_macro2::TokenStream;
use super::{ActorBuilder, BuildOptions};
fn compact_tokens(tokens: TokenStream) -> String {
tokens
.to_string()
.chars()
.filter(|character| !character.is_whitespace())
.collect()
}
fn actor_builder_with_options(options: BuildOptions) -> ActorBuilder {
ActorBuilder::new(
Arc::new(Schema::new()),
Canister::new(Def::new("test", "Canister"), "test", 0, 1, 2, 3),
options,
icydb_schema::SchemaFragment::try_new(Vec::new(), Vec::new())
.expect("empty test fragment should admit"),
)
}
#[test]
fn default_build_options_enable_minimal_metrics_only() {
let options = BuildOptions::default();
assert!(!options.sql_readonly_enabled());
assert!(!options.sql_ddl_enabled());
assert!(!options.sql_fixtures_enabled());
assert!(!options.sql_integrity_enabled());
assert!(!options.sql_introspection_enabled());
assert!(!options.sql_update_enabled());
assert_eq!(options.sql_update_policy(), None);
assert!(options.metrics_enabled());
assert!(!options.metrics_extended_enabled());
assert!(!options.snapshot_enabled());
assert!(!options.schema_enabled());
}
#[test]
fn extended_metrics_requires_metrics_surface() {
let options = BuildOptions::default()
.with_metrics_enabled(false)
.with_metrics_extended_enabled(true);
assert!(!options.metrics_enabled());
assert!(!options.metrics_extended_enabled());
let options = options.with_metrics_enabled(true);
assert!(options.metrics_enabled());
assert!(options.metrics_extended_enabled());
}
#[test]
fn generated_metrics_surface_uses_public_icydb_endpoint_names() {
let builder =
actor_builder_with_options(BuildOptions::default().with_metrics_extended_enabled(true));
let surface = compact_tokens(super::generate_metrics(&builder));
assert!(surface.contains("name=\"icydb_metrics\""));
assert!(surface.contains("name=\"icydb_metrics_extended\""));
assert!(surface.contains("name=\"icydb_metrics_reset\""));
assert!(surface.contains("pubfn__icydb_metrics("));
assert!(surface.contains("pubfn__icydb_metrics_extended("));
assert!(surface.contains("pubfn__icydb_metrics_reset("));
}
#[test]
fn generated_snapshot_surface_uses_public_icydb_endpoint_name() {
let builder =
actor_builder_with_options(BuildOptions::default().with_snapshot_enabled(true));
let surface = compact_tokens(super::generate_snapshot(&builder));
assert!(surface.contains("name=\"icydb_snapshot\""));
assert!(surface.contains("pubfn__icydb_snapshot("));
}
}