1mod accepted_schema;
7mod bounded_cache;
8mod catalog;
9mod query;
10mod response;
11#[cfg(feature = "sql")]
12mod sql;
13#[cfg(all(test, feature = "sql"))]
17mod tests;
18mod write;
19
20use crate::{
21 db::{
22 Db, EntityRuntimeHooks, FluentDeleteQuery, FluentLoadQuery, MissingRowPolicy, PersistedRow,
23 Query, StoreRegistry, WriteBatchResponse,
24 commit::CommitSchemaFingerprint,
25 executor::{DeleteExecutor, LoadExecutor, SaveExecutor},
26 schema::{AcceptedRowDecodeContract, SchemaInfo},
27 },
28 entity::{EntityKind, EntityValue},
29 error::InternalError,
30 metrics::sink::{ExecKind, MetricsSink, record_exec_error_for_path, with_metrics_sink},
31 traits::CanisterKind,
32 value::Value,
33};
34use std::thread::LocalKey;
35
36#[cfg(all(test, feature = "sql"))]
37pub(in crate::db) use accepted_schema::AcceptedCatalogRuntimeCounterSnapshot;
38pub(in crate::db) use accepted_schema::AcceptedSchemaCatalogContext;
39pub(in crate::db) use query::{
40 AcceptedExecutionOutput, AcceptedIdValuesOutput, AcceptedOptionalValueOutput,
41 AcceptedValuesOutput,
42};
43#[cfg(feature = "diagnostics")]
44pub use query::{
45 DirectDataRowAttribution, FluentTerminalExecutionAttribution, GroupedCountAttribution,
46 GroupedExecutionAttribution, KernelRowAttribution, QueryExecutionAttribution,
47 ScalarAggregateAttribution,
48};
49pub(in crate::db) use response::finalize_scalar_paged_execution;
50pub(in crate::db) use response::finalize_structural_grouped_projection_result;
51#[cfg(feature = "sql")]
52pub(in crate::db) use response::sql_grouped_cursor_from_bytes;
53#[cfg(all(feature = "sql", feature = "diagnostics"))]
54pub use sql::{
55 SqlCompileAttribution, SqlExecutionAttribution, SqlHybridCoveringAttribution,
56 SqlOutputBlobAttribution, SqlPureCoveringAttribution, SqlQueryCacheAttribution,
57 SqlQueryExecutionAttribution,
58};
59#[cfg(feature = "sql")]
60pub use sql::{
61 SqlDdlExecutionStatus, SqlDdlMutationKind, SqlDdlPreparationReport, SqlStatementDispatch,
62 SqlStatementResult, SqlStatementShellSurface, SqlStatementSurface, sql_statement_dispatch,
63 sql_statement_entity_name, sql_statement_shell_surface, sql_statement_surface,
64};
65#[cfg(all(feature = "sql", test))]
66pub(in crate::db) use sql::{
67 SqlDeleteExposurePolicy, SqlDeletePolicyContext, SqlPublicBoundedDeletePlan,
68 SqlPublicBoundedUpdatePlan, SqlPublicPrimaryKeyDeletePlan, SqlPublicPrimaryKeyUpdatePlan,
69 SqlUpdateExposurePolicy, SqlUpdatePolicyContext, SqlValidatedDeletePlan,
70 SqlValidatedUpdatePlan, classify_sql_delete_policy, classify_sql_update_policy,
71};
72#[cfg(all(feature = "sql", feature = "diagnostics"))]
73pub use sql::{SqlProjectionMaterializationMetrics, with_sql_projection_materialization_metrics};
74
75pub struct DbSession<C: CanisterKind> {
82 db: Db<C>,
83 debug: bool,
84 metrics: Option<&'static dyn MetricsSink>,
85}
86
87impl<C: CanisterKind> DbSession<C> {
88 #[must_use]
90 pub(crate) const fn new(db: Db<C>) -> Self {
91 Self {
92 db,
93 debug: false,
94 metrics: None,
95 }
96 }
97
98 #[must_use]
100 pub const fn new_with_hooks(
101 store: &'static LocalKey<StoreRegistry>,
102 entity_runtime_hooks: &'static [EntityRuntimeHooks<C>],
103 ) -> Self {
104 Self::new(Db::new_with_hooks(store, entity_runtime_hooks))
105 }
106
107 #[must_use]
109 pub const fn debug(mut self) -> Self {
110 self.debug = true;
111 self
112 }
113
114 #[must_use]
116 pub const fn metrics_sink(mut self, sink: &'static dyn MetricsSink) -> Self {
117 self.metrics = Some(sink);
118 self
119 }
120
121 const fn fluent_load_query<E>(&self, consistency: MissingRowPolicy) -> FluentLoadQuery<'_, E>
124 where
125 E: EntityKind<Canister = C>,
126 {
127 FluentLoadQuery::new(self, consistency)
128 }
129
130 fn fluent_delete_query<E>(&self, consistency: MissingRowPolicy) -> FluentDeleteQuery<'_, E>
134 where
135 E: PersistedRow<Canister = C>,
136 {
137 FluentDeleteQuery::new(self, Query::new(consistency).delete())
138 }
139
140 fn with_metrics<T>(&self, f: impl FnOnce() -> T) -> T {
141 if let Some(sink) = self.metrics {
142 with_metrics_sink(sink, f)
143 } else {
144 f()
145 }
146 }
147
148 fn execute_save_with<E, T, R>(
150 &self,
151 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
152 map: impl FnOnce(T) -> R,
153 ) -> Result<R, InternalError>
154 where
155 E: PersistedRow<Canister = C>,
156 {
157 let (contract, schema_info, schema_fingerprint) = match self
158 .with_metrics(|| self.ensure_generated_compatible_accepted_save_schema::<E>())
159 {
160 Ok(authority) => authority,
161 Err(error) => {
162 self.with_metrics(|| record_exec_error_for_path(ExecKind::Save, E::PATH, &error));
163
164 return Err(error);
165 }
166 };
167 let value = self.with_metrics(|| {
168 op(self.save_executor::<E>(contract, schema_info, schema_fingerprint))
169 })?;
170
171 Ok(map(value))
172 }
173
174 fn execute_save_with_checked_accepted_row_contract<E, T, R>(
179 &self,
180 accepted_row_decode_contract: AcceptedRowDecodeContract,
181 accepted_schema_info: SchemaInfo,
182 accepted_schema_fingerprint: CommitSchemaFingerprint,
183 op: impl FnOnce(SaveExecutor<E>) -> Result<T, InternalError>,
184 map: impl FnOnce(T) -> R,
185 ) -> Result<R, InternalError>
186 where
187 E: PersistedRow<Canister = C>,
188 {
189 let value = self.with_metrics(|| {
190 op(self.save_executor::<E>(
191 accepted_row_decode_contract,
192 accepted_schema_info,
193 accepted_schema_fingerprint,
194 ))
195 })?;
196
197 Ok(map(value))
198 }
199
200 fn execute_save_entity<E>(
202 &self,
203 op: impl FnOnce(SaveExecutor<E>) -> Result<E, InternalError>,
204 ) -> Result<E, InternalError>
205 where
206 E: PersistedRow<Canister = C>,
207 {
208 self.execute_save_with(op, std::convert::identity)
209 }
210
211 fn execute_save_batch<E>(
212 &self,
213 op: impl FnOnce(SaveExecutor<E>) -> Result<Vec<E>, InternalError>,
214 ) -> Result<WriteBatchResponse<E>, InternalError>
215 where
216 E: PersistedRow<Canister = C>,
217 {
218 self.execute_save_with(op, WriteBatchResponse::new)
219 }
220
221 #[must_use]
227 pub const fn load<E>(&self) -> FluentLoadQuery<'_, E>
228 where
229 E: EntityKind<Canister = C>,
230 {
231 self.fluent_load_query(MissingRowPolicy::Ignore)
232 }
233
234 #[must_use]
236 pub const fn load_with_consistency<E>(
237 &self,
238 consistency: MissingRowPolicy,
239 ) -> FluentLoadQuery<'_, E>
240 where
241 E: EntityKind<Canister = C>,
242 {
243 self.fluent_load_query(consistency)
244 }
245
246 #[must_use]
248 pub fn delete<E>(&self) -> FluentDeleteQuery<'_, E>
249 where
250 E: PersistedRow<Canister = C>,
251 {
252 self.fluent_delete_query(MissingRowPolicy::Ignore)
253 }
254
255 #[must_use]
257 pub fn delete_with_consistency<E>(
258 &self,
259 consistency: MissingRowPolicy,
260 ) -> FluentDeleteQuery<'_, E>
261 where
262 E: PersistedRow<Canister = C>,
263 {
264 self.fluent_delete_query(consistency)
265 }
266
267 #[must_use]
271 pub const fn select_one(&self) -> Value {
272 Value::Int64(1)
273 }
274
275 #[must_use]
280 pub(in crate::db) const fn load_executor<E>(&self) -> LoadExecutor<E>
281 where
282 E: EntityKind<Canister = C> + EntityValue,
283 {
284 LoadExecutor::new(self.db, self.debug)
285 }
286
287 #[must_use]
288 pub(in crate::db) const fn delete_executor<E>(&self) -> DeleteExecutor<E>
289 where
290 E: PersistedRow<Canister = C>,
291 {
292 DeleteExecutor::new(self.db)
293 }
294
295 #[must_use]
296 pub(in crate::db) const fn save_executor<E>(
297 &self,
298 accepted_row_decode_contract: AcceptedRowDecodeContract,
299 accepted_schema_info: SchemaInfo,
300 accepted_schema_fingerprint: CommitSchemaFingerprint,
301 ) -> SaveExecutor<E>
302 where
303 E: PersistedRow<Canister = C>,
304 {
305 SaveExecutor::new_with_accepted_contract(
306 self.db,
307 self.debug,
308 accepted_row_decode_contract,
309 accepted_schema_info,
310 accepted_schema_fingerprint,
311 )
312 }
313}