1mod query;
7#[cfg(feature = "sql")]
8mod sql;
9#[cfg(all(test, feature = "sql"))]
13mod tests;
14mod write;
15
16use crate::{
17 db::{
18 Db, EntityFieldDescription, EntitySchemaDescription, FluentDeleteQuery, FluentLoadQuery,
19 IntegrityReport, MigrationPlan, MigrationRunOutcome, MissingRowPolicy, PersistedRow, Query,
20 QueryError, StorageReport, StoreRegistry, WriteBatchResponse,
21 commit::EntityRuntimeHooks,
22 cursor::decode_optional_cursor_token,
23 executor::{DeleteExecutor, LoadExecutor, SaveExecutor},
24 schema::{describe_entity_model, show_indexes_for_model},
25 },
26 error::InternalError,
27 metrics::sink::{MetricsSink, with_metrics_sink},
28 model::entity::EntityModel,
29 traits::{CanisterKind, EntityKind, EntityValue},
30 value::Value,
31};
32use std::thread::LocalKey;
33
34#[cfg(feature = "sql")]
35pub use sql::{SqlDispatchResult, SqlParsedStatement, SqlStatementRoute};
36
37fn decode_optional_cursor_bytes(cursor_token: Option<&str>) -> Result<Option<Vec<u8>>, QueryError> {
40 decode_optional_cursor_token(cursor_token).map_err(QueryError::from_cursor_plan_error)
41}
42
43pub struct DbSession<C: CanisterKind> {
50 db: Db<C>,
51 debug: bool,
52 metrics: Option<&'static dyn MetricsSink>,
53}
54
55impl<C: CanisterKind> DbSession<C> {
56 #[must_use]
58 pub(crate) const fn new(db: Db<C>) -> Self {
59 Self {
60 db,
61 debug: false,
62 metrics: None,
63 }
64 }
65
66 #[must_use]
68 pub const fn new_with_hooks(
69 store: &'static LocalKey<StoreRegistry>,
70 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
71 ) -> Self {
72 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
73 }
74
75 #[must_use]
77 pub const fn debug(mut self) -> Self {
78 self.debug = true;
79 self
80 }
81
82 #[must_use]
84 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
85 self.metrics = Some(sink);
86 self
87 }
88
89 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
90 if let Some(sink) = self.metrics {
91 with_metrics_sink(sink, f)
92 } else {
93 f()
94 }
95 }
96
97 fn execute_save_with<E, T, R>(
99 &self,
100 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
101 map: impl FnOnce(T) -> R,
102 ) -> Result<R, InternalError>
103 where
104 E: PersistedRow<Canister = C> + EntityValue,
105 {
106 let value = self.with_metrics(|| op(self.save_executor::<E>()))?;
107
108 Ok(map(value))
109 }
110
111 fn execute_save_entity<E>(
113 &self,
114 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
115 ) -> Result<E, InternalError>
116 where
117 E: PersistedRow<Canister = C> + EntityValue,
118 {
119 self.execute_save_with(op, std::convert::identity)
120 }
121
122 fn execute_save_batch<E>(
123 &self,
124 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
125 ) -> Result<WriteBatchResponse<E>, InternalError>
126 where
127 E: PersistedRow<Canister = C> + EntityValue,
128 {
129 self.execute_save_with(op, WriteBatchResponse::new)
130 }
131
132 #[must_use]
138 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
139 where
140 E: EntityKind<Canister = C>,
141 {
142 FluentLoadQuery::new(self, Query::new(MissingRowPolicy::Ignore))
143 }
144
145 #[must_use]
147 pub const fn load_with_consistency<E>(
148 &self,
149 consistency: MissingRowPolicy,
150 ) -> FluentLoadQuery<'_, E>
151 where
152 E: EntityKind<Canister = C>,
153 {
154 FluentLoadQuery::new(self, Query::new(consistency))
155 }
156
157 #[must_use]
159 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
160 where
161 E: PersistedRow<Canister = C>,
162 {
163 FluentDeleteQuery::new(self, Query::new(MissingRowPolicy::Ignore).delete())
164 }
165
166 #[must_use]
168 pub fn delete_with_consistency<E>(
169 &self,
170 consistency: MissingRowPolicy,
171 ) -> FluentDeleteQuery<'_, E>
172 where
173 E: PersistedRow<Canister = C>,
174 {
175 FluentDeleteQuery::new(self, Query::new(consistency).delete())
176 }
177
178 #[must_use]
182 pub const fn select_one(&self) -> Value {
183 Value::Int(1)
184 }
185
186 #[must_use]
193 pub fn show_indexes<E>(&self) -> Vec<String>
194 where
195 E: EntityKind<Canister = C>,
196 {
197 self.show_indexes_for_model(E::MODEL)
198 }
199
200 #[must_use]
202 pub fn show_indexes_for_model(&self, model: &'static EntityModel) -> Vec<String> {
203 show_indexes_for_model(model)
204 }
205
206 #[must_use]
208 pub fn show_columns<E>(&self) -> Vec<EntityFieldDescription>
209 where
210 E: EntityKind<Canister = C>,
211 {
212 self.show_columns_for_model(E::MODEL)
213 }
214
215 #[must_use]
217 pub fn show_columns_for_model(
218 &self,
219 model: &'static EntityModel,
220 ) -> Vec<EntityFieldDescription> {
221 describe_entity_model(model).fields().to_vec()
222 }
223
224 #[must_use]
226 pub fn show_entities(&self) -> Vec<String> {
227 self.db.runtime_entity_names()
228 }
229
230 #[must_use]
235 pub fn describe_entity<E>(&self) -> EntitySchemaDescription
236 where
237 E: EntityKind<Canister = C>,
238 {
239 self.describe_entity_model(E::MODEL)
240 }
241
242 #[must_use]
244 pub fn describe_entity_model(&self, model: &'static EntityModel) -> EntitySchemaDescription {
245 describe_entity_model(model)
246 }
247
248 pub fn storage_report(
250 &self,
251 name_to_path: &[(&'static str, &'static str)],
252 ) -> Result<StorageReport, InternalError> {
253 self.db.storage_report(name_to_path)
254 }
255
256 pub fn integrity_report(&self) -> Result<IntegrityReport, InternalError> {
258 self.db.integrity_report()
259 }
260
261 pub fn execute_migration_plan(
266 &self,
267 plan: &MigrationPlan,
268 max_steps: usize,
269 ) -> Result<MigrationRunOutcome, InternalError> {
270 self.with_metrics(|| self.db.execute_migration_plan(plan, max_steps))
271 }
272
273 #[must_use]
278 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
279 where
280 E: EntityKind<Canister = C> + EntityValue,
281 {
282 LoadExecutor::new(self.db, self.debug)
283 }
284
285 #[must_use]
286 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
287 where
288 E: PersistedRow<Canister = C> + EntityValue,
289 {
290 DeleteExecutor::new(self.db, self.debug)
291 }
292
293 #[must_use]
294 pub(in crate::db) const fn save_executor<E>(&self) -> SaveExecutor<E>
295 where
296 E: PersistedRow<Canister = C> + EntityValue,
297 {
298 SaveExecutor::new(self.db, self.debug)
299 }
300}