1mod commit;
2pub mod executor;
3pub mod identity;
4pub mod index;
5pub mod query;
6pub mod response;
7pub mod store;
8pub mod traits;
9pub mod types;
10mod write;
11
12pub(crate) use commit::*;
13pub(crate) use write::WriteUnit;
14
15use crate::{
16 db::{
17 executor::{Context, DeleteExecutor, LoadExecutor, SaveExecutor, UpsertExecutor},
18 index::IndexStoreRegistry,
19 query::{
20 Query, QueryError,
21 diagnostics::{
22 QueryDiagnostics, QueryExecutionDiagnostics, QueryTraceExecutorKind, finish_event,
23 start_event, trace_access_from_plan,
24 },
25 },
26 response::Response,
27 store::DataStoreRegistry,
28 traits::FromKey,
29 },
30 error::InternalError,
31 obs::sink::{self, MetricsSink},
32 traits::{CanisterKind, EntityKind},
33};
34use std::{marker::PhantomData, thread::LocalKey};
35
36pub struct Db<C: CanisterKind> {
52 data: &'static LocalKey<DataStoreRegistry>,
53 index: &'static LocalKey<IndexStoreRegistry>,
54 _marker: PhantomData<C>,
55}
56
57impl<C: CanisterKind> Db<C> {
58 #[must_use]
59 pub const fn new(
60 data: &'static LocalKey<DataStoreRegistry>,
61 index: &'static LocalKey<IndexStoreRegistry>,
62 ) -> Self {
63 Self {
64 data,
65 index,
66 _marker: PhantomData,
67 }
68 }
69
70 #[must_use]
71 pub const fn context<E>(&self) -> Context<'_, E>
72 where
73 E: EntityKind<Canister = C>,
74 {
75 Context::new(self)
76 }
77
78 pub fn with_data<R>(&self, f: impl FnOnce(&DataStoreRegistry) -> R) -> R {
80 self.data.with(|reg| f(reg))
81 }
82
83 pub fn with_index<R>(&self, f: impl FnOnce(&IndexStoreRegistry) -> R) -> R {
85 self.index.with(|reg| f(reg))
86 }
87}
88
89impl<C: CanisterKind> Copy for Db<C> {}
93
94impl<C: CanisterKind> Clone for Db<C> {
95 fn clone(&self) -> Self {
96 *self
97 }
98}
99
100pub struct DbSession<C: CanisterKind> {
106 db: Db<C>,
107 debug: bool,
108 metrics: Option<&'static dyn MetricsSink>,
109}
110
111impl<C: CanisterKind> DbSession<C> {
112 #[must_use]
113 pub const fn new(db: Db<C>) -> Self {
115 Self {
116 db,
117 debug: false,
118 metrics: None,
119 }
120 }
121
122 #[must_use]
123 pub const fn debug(mut self) -> Self {
125 self.debug = true;
126 self
127 }
128
129 #[must_use]
130 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
132 self.metrics = Some(sink);
133 self
134 }
135
136 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
137 if let Some(sink) = self.metrics {
138 sink::with_metrics_sink(sink, f)
139 } else {
140 f()
141 }
142 }
143
144 #[must_use]
151 pub const fn load<E>(&self) -> LoadExecutor<E>
152 where
153 E: EntityKind<Canister = C>,
154 {
155 LoadExecutor::new(self.db, self.debug)
156 }
157
158 #[must_use]
163 pub const fn save<E>(&self) -> SaveExecutor<E>
164 where
165 E: EntityKind<Canister = C>,
166 {
167 SaveExecutor::new(self.db, self.debug)
168 }
169
170 #[must_use]
173 pub const fn upsert<E>(&self) -> UpsertExecutor<E>
174 where
175 E: EntityKind<Canister = C>,
176 E::PrimaryKey: FromKey,
177 {
178 UpsertExecutor::new(self.db, self.debug)
179 }
180
181 #[must_use]
184 pub const fn delete<E>(&self) -> DeleteExecutor<E>
185 where
186 E: EntityKind<Canister = C>,
187 {
188 DeleteExecutor::new(self.db, self.debug)
189 }
190
191 pub fn diagnose_query<E: EntityKind<Canister = C>>(
197 &self,
198 query: &Query<E>,
199 ) -> Result<QueryDiagnostics, QueryError> {
200 let explain = query.explain()?;
201 Ok(QueryDiagnostics::from(explain))
202 }
203
204 pub fn execute_with_diagnostics<E: EntityKind<Canister = C>>(
206 &self,
207 query: &Query<E>,
208 ) -> Result<(Response<E>, QueryExecutionDiagnostics), QueryError> {
209 let plan = query.plan()?;
210 let fingerprint = plan.fingerprint();
211 let access = trace_access_from_plan(plan.access());
212 let start = start_event(fingerprint, access, QueryTraceExecutorKind::Load);
213
214 let result = self.with_metrics(|| self.load::<E>().execute(plan));
215 match result {
216 Ok(response) => {
217 let rows = u64::try_from(response.0.len()).unwrap_or(u64::MAX);
218 let finish = finish_event(fingerprint, access, QueryTraceExecutorKind::Load, rows);
219 let diagnostics = QueryExecutionDiagnostics {
220 fingerprint,
221 events: vec![start, finish],
222 };
223 Ok((response, diagnostics))
224 }
225 Err(err) => Err(QueryError::Execute(err)),
226 }
227 }
228
229 pub fn insert<E>(&self, entity: E) -> Result<E, InternalError>
235 where
236 E: EntityKind<Canister = C>,
237 {
238 self.with_metrics(|| self.save::<E>().insert(entity))
239 }
240
241 pub fn insert_many<E>(
243 &self,
244 entities: impl IntoIterator<Item = E>,
245 ) -> Result<Vec<E>, InternalError>
246 where
247 E: EntityKind<Canister = C>,
248 {
249 self.with_metrics(|| self.save::<E>().insert_many(entities))
250 }
251
252 pub fn replace<E>(&self, entity: E) -> Result<E, InternalError>
254 where
255 E: EntityKind<Canister = C>,
256 {
257 self.with_metrics(|| self.save::<E>().replace(entity))
258 }
259
260 pub fn replace_many<E>(
262 &self,
263 entities: impl IntoIterator<Item = E>,
264 ) -> Result<Vec<E>, InternalError>
265 where
266 E: EntityKind<Canister = C>,
267 {
268 self.with_metrics(|| self.save::<E>().replace_many(entities))
269 }
270
271 pub fn update<E>(&self, entity: E) -> Result<E, InternalError>
273 where
274 E: EntityKind<Canister = C>,
275 {
276 self.with_metrics(|| self.save::<E>().update(entity))
277 }
278
279 pub fn update_many<E>(
281 &self,
282 entities: impl IntoIterator<Item = E>,
283 ) -> Result<Vec<E>, InternalError>
284 where
285 E: EntityKind<Canister = C>,
286 {
287 self.with_metrics(|| self.save::<E>().update_many(entities))
288 }
289
290 pub fn insert_view<E>(&self, view: E::ViewType) -> Result<E::ViewType, InternalError>
292 where
293 E: EntityKind<Canister = C>,
294 {
295 self.with_metrics(|| self.save::<E>().insert_view(view))
296 }
297
298 pub fn replace_view<E>(&self, view: E::ViewType) -> Result<E::ViewType, InternalError>
300 where
301 E: EntityKind<Canister = C>,
302 {
303 self.with_metrics(|| self.save::<E>().replace_view(view))
304 }
305
306 pub fn update_view<E>(&self, view: E::ViewType) -> Result<E::ViewType, InternalError>
308 where
309 E: EntityKind<Canister = C>,
310 {
311 self.with_metrics(|| self.save::<E>().update_view(view))
312 }
313}