1use std::sync::Arc;
5
6use runlight::BoxFuture;
7use runlight::store::{Db, DbError, Dialect, Hold, Param, Row, fill_placeholders, mysql_literal};
8use sqlx::pool::PoolConnection;
9use sqlx::{AssertSqlSafe, MySql, MySqlPool};
10use tokio::sync::Mutex;
11
12use crate::rows::{db_error, mysql_row};
13
14const LOCK: &str = "CONCAT('runlight:', LEFT(SHA2(COALESCE(DATABASE(), ''), 256), 48))";
17
18#[derive(Clone)]
22pub struct MysqlDb {
23 pool: MySqlPool,
24}
25
26impl MysqlDb {
27 pub fn new(pool: MySqlPool) -> MysqlDb {
29 MysqlDb { pool }
30 }
31}
32
33const NO_BACKSLASH_ESCAPES_OFF: &str = "SET SESSION sql_mode = TRIM(BOTH ',' FROM REPLACE(CONCAT(',', @@SESSION.sql_mode, ','), ',NO_BACKSLASH_ESCAPES,', ','))";
35
36fn text(sql: &str, params: &[Param]) -> Result<AssertSqlSafe<String>, DbError> {
40 Ok(AssertSqlSafe(format!("{NO_BACKSLASH_ESCAPES_OFF}; {}", fill_placeholders(sql, params, true, mysql_literal)?)))
41}
42
43async fn exec(conn: &mut PoolConnection<MySql>, sql: &str) -> Result<Vec<Row>, DbError> {
44 let rows = sqlx::raw_sql(AssertSqlSafe(sql.to_string())).fetch_all(&mut **conn).await.map_err(db_error)?;
45 rows.iter().map(mysql_row).collect()
46}
47
48impl Db for MysqlDb {
49 fn dialect(&self) -> Dialect {
50 Dialect::Mysql
51 }
52
53 fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
54 Box::pin(async move {
55 let rows = sqlx::raw_sql(text(sql, ¶ms)?).fetch_all(&self.pool).await.map_err(db_error)?;
56 rows.iter().map(mysql_row).collect()
57 })
58 }
59
60 fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
61 Box::pin(async move {
62 Ok(sqlx::raw_sql(text(sql, ¶ms)?).execute(&self.pool).await.map_err(db_error)?.rows_affected())
63 })
64 }
65
66 fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
67 Box::pin(async move {
68 let mut slot = Slot { conn: Some(self.pool.acquire().await.map_err(db_error)?), nested: false };
71 let conn = slot.conn.as_mut().ok_or_else(ended)?;
72 match kind {
73 Hold::Transaction => begin(conn).await?,
74 Hold::Exclusive => loop {
75 let rows = exec(conn, &format!("SELECT GET_LOCK({LOCK}, 5) AS ok")).await?;
76 let row = rows.into_iter().next().unwrap_or_default();
77 if row.is_null("ok") {
78 return Err(DbError("Runlight: MySQL refused the lock for creating tables".into()));
79 }
80 if row.num("ok") == 1.0 {
81 break;
82 }
83 },
85 }
86 Ok(Arc::new(Held { conn: Arc::new(Mutex::new(slot)), kind, role: Role::Owner }) as Arc<dyn Db>)
87 })
88 }
89}
90
91async fn begin(conn: &mut PoolConnection<MySql>) -> Result<(), DbError> {
94 exec(conn, "SET TRANSACTION ISOLATION LEVEL READ COMMITTED").await?;
95 exec(conn, "START TRANSACTION").await?;
96 Ok(())
97}
98
99struct Slot {
102 conn: Option<PoolConnection<MySql>>,
103 nested: bool,
105}
106
107impl Slot {
108 fn close(&mut self) {
109 if let Some(conn) = self.conn.take() {
110 drop(conn.detach());
111 }
112 }
113}
114
115impl Drop for Slot {
116 fn drop(&mut self) {
117 self.close();
118 }
119}
120
121struct Held {
123 conn: Arc<Mutex<Slot>>,
124 kind: Hold,
125 role: Role,
126}
127
128impl Drop for Held {
129 fn drop(&mut self) {
130 if self.role == Role::Owner
132 && let Ok(mut slot) = self.conn.try_lock()
133 {
134 slot.close();
135 }
136 }
137}
138
139#[derive(Clone, Copy, PartialEq, Eq)]
141enum Role {
142 Owner,
144 Nested,
146 Joined,
148}
149
150fn ended() -> DbError {
151 DbError("Runlight: the held connection has been let go".into())
152}
153
154impl Db for Held {
155 fn dialect(&self) -> Dialect {
156 Dialect::Mysql
157 }
158
159 fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
160 Box::pin(async move {
161 let mut slot = self.conn.lock().await;
162 let conn = slot.conn.as_mut().ok_or_else(ended)?;
163 let rows = sqlx::raw_sql(text(sql, ¶ms)?).fetch_all(&mut **conn).await.map_err(db_error)?;
164 rows.iter().map(mysql_row).collect()
165 })
166 }
167
168 fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
169 Box::pin(async move {
170 let mut slot = self.conn.lock().await;
171 let conn = slot.conn.as_mut().ok_or_else(ended)?;
172 Ok(sqlx::raw_sql(text(sql, ¶ms)?).execute(&mut **conn).await.map_err(db_error)?.rows_affected())
173 })
174 }
175
176 fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
177 Box::pin(async move {
178 if kind == Hold::Transaction && self.kind == Hold::Exclusive {
179 let mut slot = self.conn.lock().await;
180 let slot = &mut *slot;
181 let conn = slot.conn.as_mut().ok_or_else(ended)?;
182 slot.nested = true;
184 begin(conn).await?;
185 return Ok(Arc::new(Held { conn: self.conn.clone(), kind, role: Role::Nested }) as Arc<dyn Db>);
186 }
187 Ok(Arc::new(Held { conn: self.conn.clone(), kind: self.kind, role: Role::Joined }) as Arc<dyn Db>)
188 })
189 }
190
191 fn finish(&self, ok: bool) -> BoxFuture<'_, Result<(), DbError>> {
192 Box::pin(async move {
193 let mut slot = self.conn.lock().await;
194 match self.role {
195 Role::Joined => Ok(()),
196 Role::Nested => {
197 let slot = &mut *slot;
198 let Some(conn) = slot.conn.as_mut() else { return Ok(()) };
199 exec(conn, if ok { "COMMIT" } else { "ROLLBACK" }).await?;
200 slot.nested = false;
201 Ok(())
202 }
203 Role::Owner if slot.nested => {
206 slot.close();
207 Ok(())
208 }
209 Role::Owner => {
210 let Some(mut conn) = slot.conn.take() else { return Ok(()) };
211 let result = match self.kind {
212 Hold::Transaction => exec(&mut conn, if ok { "COMMIT" } else { "ROLLBACK" }).await.map(|_| ()),
213 Hold::Exclusive => exec(&mut conn, &format!("DO RELEASE_LOCK({LOCK})")).await.map(|_| ()),
216 };
217 if result.is_err() {
218 drop(conn.detach());
221 }
222 if self.kind == Hold::Exclusive { Ok(()) } else { result }
223 }
224 }
225 })
226 }
227}