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

saddle_db/
database.rs
