dbmcp_sql/
timeout.rs

1//! Query-level timeout wrapper for SQL operations.
2//!
3//! Provides [`execute_with_timeout`] which runs a query operation under
4//! an optional `tokio::time::timeout` guard.  All backend crates use this
5//! single function instead of duplicating timeout logic.
6
7use std::time::{Duration, Instant};
8
9use sqlx::SqlStr;
10
11use crate::SqlError;
12
13/// Runs a query operation with an optional query timeout.
14///
15/// `op` is handed the [`SqlStr`] to execute and returns the in-flight
16/// query future. When `timeout_secs` is `Some(n)` with `n > 0`, that
17/// future is wrapped with [`tokio::time::timeout`]; on expiry it is
18/// dropped (cancelling the in-flight query) and [`SqlError::QueryTimeout`]
19/// is returned with the elapsed time and the SQL text. `None` or
20/// `Some(0)` runs without a timeout.
21///
22/// The SQL text is only copied into the error on the timeout path, so the
23/// success path adds no allocation.
24///
25/// # Errors
26///
27/// * [`SqlError::QueryTimeout`] — the query exceeded the configured
28///   timeout.
29/// * [`SqlError::Connection`] — the database could not be reached
30///   (I/O, TLS, protocol, or connection pool failure).
31/// * [`SqlError::Query`] — the underlying query failed for any other
32///   reason (e.g. syntax error).
33pub async fn execute_with_timeout<T>(
34    timeout_secs: Option<u64>,
35    sql: SqlStr,
36    op: impl AsyncFnOnce(SqlStr) -> Result<T, sqlx::Error>,
37) -> Result<T, SqlError> {
38    let result = match timeout_secs {
39        Some(secs) if secs > 0 => {
40            let start = Instant::now();
41            let err_sql = sql.clone();
42            tokio::time::timeout(Duration::from_secs(secs), op(sql))
43                .await
44                .map_err(|_| SqlError::QueryTimeout {
45                    elapsed_secs: start.elapsed().as_secs_f64(),
46                    sql: err_sql.as_str().to_owned(),
47                })?
48        }
49        _ => op(sql).await,
50    };
51    result.map_err(|e| match e {
52        sqlx::Error::Io(_)
53        | sqlx::Error::Tls(_)
54        | sqlx::Error::Protocol(_)
55        | sqlx::Error::PoolTimedOut
56        | sqlx::Error::PoolClosed
57        | sqlx::Error::WorkerCrashed => SqlError::Connection(e.to_string()),
58        e => SqlError::Query(e.to_string()),
59    })
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[tokio::test]
67    async fn fast_query_succeeds_with_timeout() {
68        let result = execute_with_timeout(Some(5), SqlStr::from_static("SELECT 1"), |_| async { Ok(42) }).await;
69        assert_eq!(result.expect("should succeed"), 42);
70    }
71
72    #[tokio::test]
73    async fn query_error_propagates_as_app_error() {
74        let result: Result<i32, SqlError> = execute_with_timeout(Some(5), SqlStr::from_static("BAD SQL"), |_| async {
75            Err(sqlx::Error::Configuration("syntax error".into()))
76        })
77        .await;
78        let err = result.expect_err("should fail");
79        assert!(
80            matches!(err, SqlError::Query(ref msg) if msg.contains("syntax error")),
81            "unexpected error: {err}"
82        );
83    }
84
85    #[tokio::test]
86    async fn connectivity_errors_become_connection_errors() {
87        for error in [
88            sqlx::Error::Io(std::io::Error::new(std::io::ErrorKind::ConnectionReset, "reset")),
89            sqlx::Error::Tls("bad certificate".into()),
90            sqlx::Error::Protocol("unexpected eof".into()),
91            sqlx::Error::PoolTimedOut,
92            sqlx::Error::PoolClosed,
93            sqlx::Error::WorkerCrashed,
94        ] {
95            let message = error.to_string();
96            let result: Result<i32, SqlError> =
97                execute_with_timeout(None, SqlStr::from_static("SELECT 1"), |_| async { Err(error) }).await;
98            let err = result.expect_err("should fail");
99            assert!(matches!(err, SqlError::Connection(ref m) if *m == message), "{err:?}");
100            assert_eq!(err.to_string(), format!("Database error: {message}"));
101        }
102    }
103
104    #[tokio::test]
105    async fn slow_query_times_out() {
106        let result: Result<i32, SqlError> =
107            execute_with_timeout(Some(1), SqlStr::from_static("SELECT SLEEP(60)"), |_| async {
108                tokio::time::sleep(Duration::from_mins(1)).await;
109                Ok(0)
110            })
111            .await;
112        let err = result.expect_err("should time out");
113        match err {
114            SqlError::QueryTimeout { elapsed_secs, sql } => {
115                assert!(elapsed_secs >= 0.9, "elapsed too small: {elapsed_secs}");
116                assert_eq!(sql, "SELECT SLEEP(60)");
117            }
118            other => panic!("expected QueryTimeout, got: {other}"),
119        }
120    }
121
122    #[tokio::test]
123    async fn none_timeout_runs_without_limit() {
124        let result = execute_with_timeout(None, SqlStr::from_static("SELECT 1"), |_| async { Ok(1) }).await;
125        assert_eq!(result.expect("should succeed"), 1);
126    }
127
128    #[tokio::test]
129    async fn zero_timeout_disables_limit() {
130        let result = execute_with_timeout(Some(0), SqlStr::from_static("SELECT 1"), |_| async { Ok(1) }).await;
131        assert_eq!(result.expect("should succeed"), 1);
132    }
133}
dbmcp_sql/timeout.rs

dbmcp_sql/
timeout.rs
