ocre/runtime/ctx.rs
1use std::{
2 collections::HashMap,
3 sync::{Arc, Mutex, PoisonError},
4};
5
6use axum::http::{HeaderMap, Method, header};
7use serde::de::DeserializeOwned;
8use worker::{
9 D1Database, D1DatabaseSession, Env,
10 js_sys::Array,
11 send::{SendFuture, SendWrapper},
12};
13
14use super::{Db, d1::Handle};
15use crate::{
16 Error, Result, cache::QUERY_CACHE_LIMIT, errors::Reporter, events::Events, instrument::Timings, log::Logger,
17};
18
19/// Name of the D1 binding every Ocre app uses for its main database.
20const DB_BINDING: &str = "DB";
21
22/// Per-request application context: the Worker environment with typed access to its bindings.
23///
24/// [`serve`](crate::serve) builds one per request and passes it to the router
25/// as axum state, so handlers take `State(ctx): State<Ctx>`. Queue consumers,
26/// cron runs and inbound email handlers receive one too. Cloning is cheap (the
27/// environment is a JavaScript handle) and the type is `Send`, so it can be
28/// held across `.await` in plain axum handlers.
29///
30/// Creating a `Ctx` costs nothing on the free plan; each binding call
31/// ([`db`](Self::db), `ocre::cache`, `ocre::storage`...) only looks the
32/// binding up.
33///
34/// A `Ctx` and its clones share two per-request memories, like Rails' query
35/// cache and local cache: identical `SELECT`s through [`db`](Self::db) are
36/// answered from memory until a write (see [`Db`]), and `ocre::cache`
37/// reads each KV key at most once. Each request, queued job and cron run
38/// gets its own.
39///
40/// # Examples
41///
42/// ```no_run
43/// use axum::{Router, extract::State, routing::get};
44/// use ocre::{ApiResult, Ctx};
45///
46/// async fn count(State(ctx): State<Ctx>) -> ApiResult<String> {
47/// let db = ctx.db()?;
48/// let rows: Vec<serde_json::Value> = db.all("SELECT COUNT(*) AS n FROM posts", ocre::params![]).await?;
49/// Ok(rows[0]["n"].to_string())
50/// }
51///
52/// fn routes() -> Router<Ctx> {
53/// Router::new().route("/count", get(count))
54/// }
55/// # let _ = routes;
56/// ```
57#[derive(Clone)]
58pub struct Ctx {
59 env: SendWrapper<Env>,
60 memo: Arc<Memo>,
61 log: Logger,
62 errors: Reporter,
63 events: Events,
64 timings: Timings,
65}
66
67/// What one request remembers: KV texts read or written (the local cache),
68/// `SELECT` results (the query cache) and, with read replicas on, its D1
69/// sessions.
70#[derive(Default)]
71pub(crate) struct Memo {
72 kv: Mutex<HashMap<String, Option<String>>>,
73 queries: Mutex<HashMap<String, SendWrapper<Array>>>,
74 replicas: Mutex<Option<Replicas>>,
75}
76
77/// The D1 sessions of a request ([`crate::replicas`]).
78struct Replicas {
79 method: Method,
80 /// The request's `Cookie` headers, holding the visitor's bookmarks.
81 cookies: HeaderMap,
82 /// One session per binding, opened on first use.
83 sessions: Vec<(String, Arc<D1DatabaseSession>)>,
84 /// Bindings that ran a statement that may write.
85 wrote: Vec<String>,
86}
87
88impl Memo {
89 /// The text of a KV key this request already read or wrote.
90 pub(crate) fn kv(&self, key: &str) -> Option<Option<String>> {
91 self.kv.lock().unwrap_or_else(PoisonError::into_inner).get(key).cloned()
92 }
93
94 pub(crate) fn remember_kv(&self, key: &str, text: Option<String>) {
95 self.kv.lock().unwrap_or_else(PoisonError::into_inner).insert(key.to_owned(), text);
96 }
97
98 /// The rows of a `SELECT` this request already ran.
99 pub(crate) fn query(&self, key: &str) -> Option<Array> {
100 self.queries.lock().unwrap_or_else(PoisonError::into_inner).get(key).map(|rows| Array::clone(rows))
101 }
102
103 pub(crate) fn remember_query(&self, key: String, rows: Array) {
104 let mut queries = self.queries.lock().unwrap_or_else(PoisonError::into_inner);
105 if queries.len() >= QUERY_CACHE_LIMIT {
106 queries.clear();
107 }
108 queries.insert(key, SendWrapper::new(rows));
109 }
110
111 /// Forgets every `SELECT` result, after a statement that may write.
112 pub(crate) fn clear_queries(&self) {
113 self.queries.lock().unwrap_or_else(PoisonError::into_inner).clear();
114 }
115
116 /// Queries of this request go through D1 sessions ([`crate::replicas`]).
117 pub(crate) fn start_replicas(&self, method: &Method, headers: &HeaderMap) {
118 let mut cookies = HeaderMap::new();
119 for value in headers.get_all(header::COOKIE) {
120 cookies.append(header::COOKIE, value.clone());
121 }
122 let replicas = Replicas { method: method.clone(), cookies, sessions: Vec::new(), wrote: Vec::new() };
123 *self.replicas.lock().unwrap_or_else(PoisonError::into_inner) = Some(replicas);
124 }
125
126 /// The request's session on `binding` (opened on first use), or `None`
127 /// when replicas are off: then queries go to the database itself.
128 fn session(&self, binding: &str, db: &D1Database) -> Option<worker::Result<Arc<D1DatabaseSession>>> {
129 let mut replicas = self.replicas.lock().unwrap_or_else(PoisonError::into_inner);
130 let replicas = replicas.as_mut()?;
131 if let Some((_, session)) = replicas.sessions.iter().find(|(name, _)| name == binding) {
132 return Some(Ok(Arc::clone(session)));
133 }
134 let start = crate::replicas::session_start(&replicas.cookies, &replicas.method, binding);
135 Some(db.with_session(Some(&start)).map(|session| {
136 let session = Arc::new(session);
137 replicas.sessions.push((binding.to_owned(), Arc::clone(&session)));
138 session
139 }))
140 }
141
142 /// Notes that `binding` ran a statement that may write.
143 pub(crate) fn wrote(&self, binding: &str) {
144 if let Some(replicas) = self.replicas.lock().unwrap_or_else(PoisonError::into_inner).as_mut()
145 && !replicas.wrote.iter().any(|name| name == binding)
146 {
147 replicas.wrote.push(binding.to_owned());
148 }
149 }
150
151 /// The bookmark of each session that wrote, for the response's cookies.
152 pub(crate) fn bookmarks(&self) -> Vec<(String, String)> {
153 let replicas = self.replicas.lock().unwrap_or_else(PoisonError::into_inner);
154 let Some(replicas) = replicas.as_ref() else { return Vec::new() };
155 let wrote = replicas.sessions.iter().filter(|(name, _)| replicas.wrote.contains(name));
156 wrote.filter_map(|(name, session)| Some((name.clone(), session.get_bookmark().ok()??))).collect()
157 }
158}
159
160impl Ctx {
161 pub(crate) fn new(env: Env) -> Self {
162 // Once per Worker instance: the keys of encrypted model columns.
163 crate::encryption::ensure_installed(&|name| env.secret(name).ok().map(|secret| secret.to_string()));
164 let var = |name: &str| env.var(name).ok().map(|var| var.to_string());
165 crate::log::configure(var(crate::log::LOG_LEVEL).as_deref(), var(crate::log::LOG_FORMAT).as_deref());
166 install_panic_hook();
167 let log = Logger::new();
168 Self {
169 env: SendWrapper::new(env),
170 memo: Arc::default(),
171 errors: Reporter::new(log.clone()),
172 events: Events::new(log.clone()),
173 log,
174 timings: Timings::default(),
175 }
176 }
177
178 /// The same context whose log lines and error reports carry `log`'s fields (the request id...).
179 pub(crate) fn with_log(mut self, log: Logger) -> Self {
180 self.errors = Reporter::new(log.clone());
181 self.events = Events::new(log.clone());
182 self.log = log;
183 self
184 }
185
186 /// The same environment with empty per-request memories, e.g. for each job of a queue batch.
187 pub(crate) fn fresh(&self) -> Self {
188 Self { memo: Arc::default(), ..self.clone() }
189 }
190
191 pub(crate) fn memo(&self) -> &Arc<Memo> {
192 &self.memo
193 }
194
195 pub(crate) fn timings(&self) -> &Timings {
196 &self.timings
197 }
198
199 /// The logger of this request, job batch or cron run (see [`ocre::log`](crate::log)).
200 ///
201 /// In a request its lines carry `request_id`, `method` and `path`, the
202 /// same request id as the [`RequestId`](crate::RequestId) extractor and
203 /// the `X-Request-Id` response header. Add fields with
204 /// [`Logger::with`]. No binding call; each line is one Workers Logs event.
205 ///
206 /// # Examples
207 ///
208 /// ```no_run
209 /// use axum::extract::State;
210 /// use ocre::Ctx;
211 ///
212 /// async fn import(State(ctx): State<Ctx>) -> &'static str {
213 /// let log = ctx.log().with("import_id", 12);
214 /// log.info("import started");
215 /// log.warn(format_args!("{} rows skipped", 3));
216 /// "OK"
217 /// }
218 /// # let _ = import;
219 /// ```
220 pub fn log(&self) -> &Logger {
221 &self.log
222 }
223
224 /// The error reporter of this request, job batch or cron run (see [`ocre::errors`](crate::errors)).
225 ///
226 /// Reports are logged at once and sent to the registered subscribers
227 /// when the response is ready (or the job or cron run ends).
228 ///
229 /// # Examples
230 ///
231 /// ```no_run
232 /// use axum::extract::State;
233 /// use ocre::{Ctx, Result};
234 ///
235 /// async fn refresh(State(ctx): State<Ctx>) -> Result<&'static str> {
236 /// ctx.errors().set_context("feed", "rss");
237 /// let body: Option<String> = ctx.errors().handle("<rss/>".parse::<String>());
238 /// Ok(if body.is_some() { "refreshed" } else { "kept the old feed" })
239 /// }
240 /// # let _ = refresh;
241 /// ```
242 pub fn errors(&self) -> &Reporter {
243 &self.errors
244 }
245
246 /// Structured events of this request or job (Rails' `Rails.event`), see [`ocre::events`](crate::events).
247 ///
248 /// # Examples
249 ///
250 /// ```no_run
251 /// use axum::extract::State;
252 /// use ocre::Ctx;
253 ///
254 /// async fn signup(State(ctx): State<Ctx>) -> &'static str {
255 /// ctx.events().notify("user.signed_up", serde_json::json!({ "plan": "free" }));
256 /// "welcome"
257 /// }
258 /// # let _ = signup;
259 /// ```
260 pub fn events(&self) -> &Events {
261 &self.events
262 }
263
264 /// The app's settings, read from Worker variables and secrets into `T` (see [`ocre::config`](crate::config)).
265 ///
266 /// Each field reads the variable of the same name in upper case, or
267 /// else the secret. No binding call: one environment lookup per field.
268 ///
269 /// # Errors
270 ///
271 /// [`Error::Internal`] (500, logged) naming the variable when a required
272 /// one is missing or does not convert to its field's type.
273 ///
274 /// # Examples
275 ///
276 /// ```no_run
277 /// use axum::extract::State;
278 /// use ocre::{Ctx, Result};
279 /// use serde::Deserialize;
280 ///
281 /// #[derive(Deserialize)]
282 /// struct Settings {
283 /// support_email: String,
284 /// max_uploads: Option<u32>,
285 /// }
286 ///
287 /// async fn contact(State(ctx): State<Ctx>) -> Result<String> {
288 /// let settings: Settings = ctx.config()?;
289 /// Ok(settings.support_email)
290 /// }
291 /// # let _ = contact;
292 /// ```
293 pub fn config<T: DeserializeOwned>(&self) -> Result<T> {
294 crate::config::from_lookup(&|name| super::errors::lookup(&self.env, name))
295 }
296
297 /// The raw Workers environment, for bindings Ocre does not wrap yet.
298 ///
299 /// Use it for vars, secrets and bindings such as AI or Vectorize; prefer
300 /// the Ocre helpers when one exists, since their errors name the
301 /// cloudflare.config.ts fix.
302 ///
303 /// # Examples
304 ///
305 /// ```no_run
306 /// use axum::extract::State;
307 /// use ocre::{Ctx, Result};
308 ///
309 /// async fn app_name(State(ctx): State<Ctx>) -> Result<String> {
310 /// let name = ctx.env().var("APP_NAME")?;
311 /// Ok(name.to_string())
312 /// }
313 /// # let _ = app_name;
314 /// ```
315 pub fn env(&self) -> &Env {
316 &self.env
317 }
318
319 /// The secret `name`, wherever it is kept: a Worker secret (`ocre secrets
320 /// push`), a `.dev.vars` value in `ocre dev`, or a secret of the
321 /// account's Secrets Store bound to the Worker in cloudflare.config.ts
322 /// (`NAME: bindings.secretsStoreSecret({ storeId, secretName })`, which
323 /// `ocre secrets push NAME --store` writes). Code reads it the same way
324 /// wherever it lives, so moving a secret to the store changes no code.
325 ///
326 /// A Worker secret is read without I/O; a Secrets Store secret costs one
327 /// call to the store. Secrets Ocre reads without awaiting
328 /// (`SECRET_KEY_BASE`, `R2_*`) must stay Worker secrets.
329 ///
330 /// # Errors
331 ///
332 /// [`Error::Internal`] (500, logged) when the secret is missing or
333 /// empty, naming where to set it.
334 ///
335 /// # Examples
336 ///
337 /// ```no_run
338 /// use axum::extract::State;
339 /// use ocre::{Ctx, Result};
340 ///
341 /// async fn rates(State(ctx): State<Ctx>) -> Result<String> {
342 /// let key = ctx.secret("RATES_API_KEY").await?;
343 /// Ok(format!("{} characters", key.len()))
344 /// }
345 /// # let _ = rates;
346 /// ```
347 pub fn secret(&self, name: &str) -> impl Future<Output = Result<String>> + Send + use<> {
348 let env = self.env.clone();
349 let name = name.to_owned();
350 SendFuture::new(async move { super::secrets::require(&env, &name).await })
351 }
352
353 /// The application database: the D1 binding `DB`.
354 ///
355 /// Looking the binding up runs no query and costs no D1 rows; see [`Db`]
356 /// for what each query reads and writes.
357 ///
358 /// # Errors
359 ///
360 /// [`Error::Internal`] (500, logged) when the Worker has no `DB` binding;
361 /// the message says to add `DB: bindings.d1({ name: "<app>" })` to
362 /// cloudflare.config.ts.
363 ///
364 /// # Examples
365 ///
366 /// ```no_run
367 /// use axum::extract::State;
368 /// use ocre::{Ctx, Result};
369 ///
370 /// async fn handler(State(ctx): State<Ctx>) -> Result<String> {
371 /// let db = ctx.db()?;
372 /// let changed = db.execute("DELETE FROM sessions WHERE expires_at < ?1", ocre::params![ocre::now()]).await?;
373 /// Ok(format!("{changed} expired"))
374 /// }
375 /// # let _ = handler;
376 /// ```
377 pub fn db(&self) -> Result<Db> {
378 self.db_named(DB_BINDING)
379 }
380
381 /// Another D1 database of the app, by its binding name (Rails' multiple databases).
382 ///
383 /// Each database is a `KEY: bindings.d1({ name: "..." })` entry in
384 /// cloudflare.config.ts with its own binding key and database name (see the Models
385 /// guide, "Several databases"). Queries cannot join across databases:
386 /// load ids from one, then `find_many` in the other.
387 ///
388 /// # Free plan
389 ///
390 /// Up to 10 databases per account, 5 GB of storage and the daily row
391 /// quotas shared by all of them.
392 ///
393 /// # Errors
394 ///
395 /// [`Error::Internal`] (500, logged) when the Worker has no D1 binding
396 /// named `binding`; the message names the cloudflare.config.ts entry to add.
397 ///
398 /// # Examples
399 ///
400 /// ```no_run
401 /// use axum::extract::State;
402 /// use ocre::{Ctx, Result, params};
403 ///
404 /// async fn track(State(ctx): State<Ctx>) -> Result<()> {
405 /// let analytics = ctx.db_named("ANALYTICS")?;
406 /// analytics.execute("INSERT INTO page_views (path) VALUES (?1)", params!["/"]).await?;
407 /// Ok(())
408 /// }
409 /// # let _ = track;
410 /// ```
411 pub fn db_named(&self, binding: &str) -> Result<Db> {
412 let handle = self.env.d1(binding).and_then(|db| match self.memo.session(binding, &db) {
413 Some(session) => session.map(Handle::Session),
414 None => Ok(Handle::Database(Arc::new(db))),
415 });
416 let db = handle.map(|handle| Db::new(handle, binding, Some(Arc::clone(&self.memo))));
417 db.map(|db| db.probed(self.log.clone(), self.timings.clone())).map_err(|err| {
418 Error::internal(format!(
419 "D1 binding `{binding}` is missing ({err}). Fix: add `{binding}: bindings.d1({{ name: \"<database>\" }}),` to worker.env in cloudflare.config.ts"
420 ))
421 })
422 }
423}
424
425/// Logs panics as `error` lines with their location: on Workers a panic
426/// otherwise ends the request with only `RuntimeError: unreachable`.
427fn install_panic_hook() {
428 static ONCE: std::sync::Once = std::sync::Once::new();
429 ONCE.call_once(|| {
430 std::panic::set_hook(Box::new(|info| {
431 let payload = info.payload();
432 let message = payload
433 .downcast_ref::<&str>()
434 .copied()
435 .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
436 .unwrap_or("Box<dyn Any>");
437 let location = info.location().map(|at| (at.file(), at.line(), at.column()));
438 Logger::new().error(crate::log::panic_line(message, location));
439 }));
440 });
441}