saddle_db/
database.rs

1use std::{str::FromStr, sync::Arc, time::Duration};
2
3use futures_util::TryStreamExt;
4use saddle_core::{ComponentLifecycle, LifecycleFuture, OperationId, Result};
5use saddle_observability::{CallKind, Observer};
6use sqlx::{
7    Connection, MySqlPool,
8    mysql::{MySqlConnectOptions, MySqlConnection, MySqlPoolOptions},
9};
10
11use crate::{
12    CallContext, DbRow, MAX_RESULT_BYTES, Statement, Transaction, TransactionFuture, WriteResult,
13    cleanup::CleanupCoordinator,
14    error::{invalid_config, map_operation_error, result_limit_exceeded, transaction_begin_failed},
15    row::row_payload_bytes,
16};
17
18pub const MAX_QUERY_ROWS: usize = 10_000;
19/// Maximum MySQL protocol packet accepted by the V1 deployment contract.
20///
21/// Pool startup rejects servers configured with a larger `max_allowed_packet`,
22/// and every physical pool connection repeats the check. This value remains
23/// below MySQL's `0xFF_FF_FF` continuation threshold, so sqlx receives one
24/// bounded buffer and cannot enter its two-fragment aggregate preallocation.
25pub const MAX_INBOUND_PACKET_BYTES: u64 = 8_388_608;
26
27/// Configuration for the single V1 MySQL/MariaDB data source.
28#[derive(Clone)]
29pub struct DatabaseConfig {
30    url: String,
31    max_connections: u32,
32    acquire_timeout: Duration,
33}
34
35impl DatabaseConfig {
36    pub fn new(url: impl Into<String>) -> Self {
37        Self {
38            url: url.into(),
39            max_connections: 16,
40            acquire_timeout: Duration::from_secs(5),
41        }
42    }
43    pub fn max_connections(mut self, value: u32) -> Self {
44        self.max_connections = value;
45        self
46    }
47    pub fn acquire_timeout(mut self, value: Duration) -> Self {
48        self.acquire_timeout = value;
49        self
50    }
51
52    fn options(&self) -> Result<MySqlConnectOptions> {
53        if self.max_connections == 0 {
54            return Err(invalid_config("max connections must be greater than zero"));
55        }
56        if self.acquire_timeout.is_zero() {
57            return Err(invalid_config("acquire timeout must be greater than zero"));
58        }
59        MySqlConnectOptions::from_str(&self.url)
60            .map_err(|_| invalid_config("database URL is not a valid MySQL/MariaDB URL"))
61    }
62}
63
64/// The process-wide managed database capability.
65#[derive(Clone)]
66pub struct Database {
67    pool: MySqlPool,
68    observer: Observer,
69    cleanup: Arc<CleanupCoordinator>,
70}
71
72impl Database {
73    /// Creates the single managed pool and verifies that it can connect.
74    pub async fn connect(config: DatabaseConfig, observer: Observer) -> Result<Self> {
75        let options = config.options()?;
76        let mut preflight = MySqlConnection::connect_with(&options)
77            .await
78            .map_err(map_operation_error)?;
79        let server_packet_limit = sqlx::query_scalar::<_, u64>("SELECT @@max_allowed_packet")
80            .fetch_one(&mut preflight)
81            .await
82            .map_err(map_operation_error)?;
83        preflight.close().await.map_err(map_operation_error)?;
84        if server_packet_limit > MAX_INBOUND_PACKET_BYTES {
85            return Err(invalid_config(
86                "server max_allowed_packet exceeds the V1 inbound allocation limit",
87            ));
88        }
89        let pool = MySqlPoolOptions::new()
90            .max_connections(config.max_connections)
91            .acquire_timeout(config.acquire_timeout)
92            .after_connect(|connection, _metadata| {
93                Box::pin(async move {
94                    let packet_limit = sqlx::query_scalar::<_, u64>("SELECT @@max_allowed_packet")
95                        .fetch_one(connection)
96                        .await?;
97                    if packet_limit > MAX_INBOUND_PACKET_BYTES {
98                        return Err(sqlx::Error::Protocol(
99                            "server packet limit exceeds Saddle V1 allocation boundary".to_owned(),
100                        ));
101                    }
102                    Ok(())
103                })
104            })
105            .connect_with(options)
106            .await
107            .map_err(map_operation_error)?;
108        Ok(Self {
109            pool,
110            observer,
111            cleanup: CleanupCoordinator::start(),
112        })
113    }
114
115    pub async fn query_all(
116        &self,
117        parent: &CallContext,
118        statement: Statement,
119    ) -> Result<Vec<DbRow>> {
120        statement.validate()?;
121        let operation = statement.operation().to_owned();
122        let call = self.observer.start_child_call(
123            parent,
124            CallKind::Database,
125            "database",
126            "database",
127            OperationId::from(operation),
128        );
129        let result = async {
130            let mut stream = statement.query().fetch(&self.pool);
131            let mut rows = Vec::new();
132            let mut result_bytes = 0_usize;
133            while let Some(row) = stream.try_next().await.map_err(map_operation_error)? {
134                if rows.len() == MAX_QUERY_ROWS {
135                    return Err(result_limit_exceeded());
136                }
137                result_bytes = result_bytes.saturating_add(row_payload_bytes(&row)?);
138                if result_bytes > MAX_RESULT_BYTES {
139                    return Err(result_limit_exceeded());
140                }
141                rows.push(DbRow(row));
142            }
143            Ok(rows)
144        }
145        .await;
146        finish_call(call, &result);
147        result
148    }
149
150    pub async fn query_optional(
151        &self,
152        parent: &CallContext,
153        statement: Statement,
154    ) -> Result<Option<DbRow>> {
155        statement.validate()?;
156        let operation = statement.operation().to_owned();
157        let call = self.observer.start_child_call(
158            parent,
159            CallKind::Database,
160            "database",
161            "database",
162            OperationId::from(operation),
163        );
164        let result = statement
165            .query()
166            .fetch_optional(&self.pool)
167            .await
168            .map_err(map_operation_error)
169            .and_then(|row| {
170                row.map(|row| {
171                    row_payload_bytes(&row)?;
172                    Ok(DbRow(row))
173                })
174                .transpose()
175            });
176        finish_call(call, &result);
177        result
178    }
179
180    pub async fn write(&self, parent: &CallContext, statement: Statement) -> Result<WriteResult> {
181        statement.validate()?;
182        let operation = statement.operation().to_owned();
183        let call = self.observer.start_child_call(
184            parent,
185            CallKind::Database,
186            "database",
187            "database",
188            OperationId::from(operation),
189        );
190        let result = statement
191            .query()
192            .execute(&self.pool)
193            .await
194            .map(|result| WriteResult::new(result.rows_affected(), result.last_insert_id()))
195            .map_err(map_operation_error);
196        finish_call(call, &result);
197        result
198    }
199
200    /// Runs one explicit, single-level transaction.
201    ///
202    /// The callback returns a boxed async block because that is the minimal
203    /// Rust API that safely ties all operations to the borrowed transaction.
204    pub async fn transaction<T, F>(
205        &self,
206        parent: &CallContext,
207        operation: impl Into<OperationId>,
208        work: F,
209    ) -> Result<T>
210    where
211        T: Send,
212        F: for<'a> FnOnce(&'a mut Transaction) -> TransactionFuture<'a, T> + Send,
213    {
214        let operation = operation.into();
215        crate::statement::validate_operation(operation.as_str())?;
216        let call = self.observer.start_child_call(
217            parent,
218            CallKind::Transaction,
219            "database",
220            "database",
221            operation,
222        );
223        let cleanup = match self.cleanup.transaction_sender() {
224            Some(cleanup) => cleanup,
225            None => {
226                let error = transaction_begin_failed();
227                call.fail(&error);
228                return Err(error);
229            }
230        };
231        let begin = self.observer.start_child_call(
232            call.context(),
233            CallKind::Transaction,
234            "database",
235            "database",
236            "begin",
237        );
238        let raw = match self.pool.begin().await {
239            Ok(transaction) => {
240                begin.succeed();
241                transaction
242            }
243            Err(_) => {
244                let error = transaction_begin_failed();
245                begin.fail(&error);
246                call.fail(&error);
247                return Err(error);
248            }
249        };
250        let mut transaction = Transaction::new(raw, self.observer.clone(), call, cleanup);
251        match work(&mut transaction).await {
252            Ok(value) => transaction.commit().await.map(|()| value),
253            Err(work_error) => Err(transaction.rollback(work_error).await),
254        }
255    }
256
257    async fn close(&self) -> Result<()> {
258        let cleanup = self.cleanup.shutdown().await;
259        self.pool.close().await;
260        cleanup
261    }
262
263    #[cfg(test)]
264    pub(crate) fn connect_lazy(config: DatabaseConfig, observer: Observer) -> Result<Self> {
265        let options = config.options()?;
266        let pool = MySqlPoolOptions::new()
267            .max_connections(config.max_connections)
268            .acquire_timeout(config.acquire_timeout)
269            .connect_lazy_with(options);
270        Ok(Self {
271            pool,
272            observer,
273            cleanup: CleanupCoordinator::start(),
274        })
275    }
276}
277
278impl ComponentLifecycle for Database {
279    fn name(&self) -> &'static str {
280        "database"
281    }
282    fn start(&self) -> LifecycleFuture<'_> {
283        Box::pin(async { Ok(()) })
284    }
285    fn shutdown(&self) -> LifecycleFuture<'_> {
286        Box::pin(async move { self.close().await })
287    }
288}
289
290fn finish_call<T>(call: saddle_observability::ActiveCall, result: &Result<T>) {
291    match result {
292        Ok(_) => call.succeed(),
293        Err(error) => call.fail(error),
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use saddle_core::{ApplicationId, ErrorKind, ModuleId, ServiceId, SpanId, TraceId};
300    use saddle_observability::ObserverConfig;
301    use serde_json::Value;
302    use std::{
303        io,
304        sync::{Arc, Mutex},
305    };
306
307    use super::*;
308    use crate::SaddleError;
309
310    #[derive(Clone, Default)]
311    struct Capture(Arc<Mutex<Vec<u8>>>);
312    impl io::Write for Capture {
313        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
314            self.0.lock().unwrap().extend_from_slice(bytes);
315            Ok(bytes.len())
316        }
317        fn flush(&mut self) -> io::Result<()> {
318            Ok(())
319        }
320    }
321    fn context() -> CallContext {
322        CallContext::new(
323            ApplicationId::from("shop"),
324            ModuleId::from("orders"),
325            ServiceId::from("orders"),
326            OperationId::from("create"),
327            TraceId::from_u128(1),
328            SpanId::from_u64(2),
329        )
330    }
331
332    #[test]
333    fn configuration_rejects_invalid_bounds_without_exposing_url() {
334        let observer = Observer::with_writer(ObserverConfig::default(), io::sink()).unwrap();
335        let error = Database::connect_lazy(
336            DatabaseConfig::new("mysql://secret@localhost/db").max_connections(0),
337            observer,
338        )
339        .err()
340        .unwrap();
341        assert_eq!(error.code(), "db.invalid_config");
342        assert!(!error.to_string().contains("secret"));
343    }
344
345    #[tokio::test]
346    async fn closed_pool_query_has_stable_error_and_trace_record() {
347        let capture = Capture::default();
348        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
349        let database = Database::connect_lazy(
350            DatabaseConfig::new("mysql://localhost/db"),
351            observer.clone(),
352        )
353        .unwrap();
354        database.close().await.unwrap();
355        let error = match database
356            .query_all(
357                &context(),
358                Statement::new("orders.list", "SELECT 1").unwrap(),
359            )
360            .await
361        {
362            Ok(_) => panic!("closed pool query unexpectedly succeeded"),
363            Err(error) => error,
364        };
365        assert_eq!(error.kind(), ErrorKind::Unavailable);
366        assert_eq!(error.code(), "db.connection_unavailable");
367        observer.flush().await.unwrap();
368        let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
369        let records: Vec<Value> = output
370            .lines()
371            .map(|line| serde_json::from_str(line).unwrap())
372            .collect();
373        assert_eq!(records[0]["call_kind"], "database");
374        assert_eq!(records[1]["error_code"], "db.connection_unavailable");
375        assert_eq!(records[0]["trace_id"], context().trace_id().to_string());
376        assert!(!output.contains("SELECT 1"));
377    }
378
379    #[tokio::test]
380    async fn transaction_begin_failure_records_phase_and_stable_error() {
381        let capture = Capture::default();
382        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
383        let database = Database::connect_lazy(
384            DatabaseConfig::new("mysql://localhost/db"),
385            observer.clone(),
386        )
387        .unwrap();
388        database.pool.close().await;
389        let error = database
390            .transaction(&context(), "orders.create", |_transaction| {
391                Box::pin(async { Ok::<_, SaddleError>(()) })
392            })
393            .await
394            .unwrap_err();
395        assert_eq!(error.code(), "db.transaction_begin_failed");
396        observer.flush().await.unwrap();
397        let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
398        let records: Vec<Value> = output
399            .lines()
400            .map(|line| serde_json::from_str(line).unwrap())
401            .collect();
402        assert_eq!(records.len(), 4);
403        assert_eq!(records[0]["operation"], "orders.create");
404        assert_eq!(records[1]["operation"], "begin");
405        assert_eq!(records[2]["error_code"], "db.transaction_begin_failed");
406        assert_eq!(records[3]["error_code"], "db.transaction_begin_failed");
407        assert!(
408            records
409                .iter()
410                .all(|record| record["trace_id"] == context().trace_id().to_string())
411        );
412    }
413
414    #[allow(dead_code)]
415    async fn transaction_usage_compiles(database: &Database, context: &CallContext) -> Result<()> {
416        database
417            .transaction(context, "orders.create", |transaction| {
418                Box::pin(async move {
419                    transaction
420                        .write(
421                            Statement::new("orders.insert", "INSERT INTO orders(id) VALUES (?)")?
422                                .bind(1_u64)?,
423                        )
424                        .await?;
425                    transaction
426                        .query_optional(
427                            Statement::new("orders.find", "SELECT id FROM orders WHERE id = ?")?
428                                .bind(1_u64)?,
429                        )
430                        .await?;
431                    Ok(())
432                })
433            })
434            .await
435    }
436}
saddle_db/database.rs

saddle_db/
database.rs
