tank_sqlite/
connection.rs

1use crate::{
2    CBox, SQLiteDriver, SQLitePrepared, SQLiteTransaction,
3    extract::{extract_name, extract_value},
4};
5use async_stream::try_stream;
6use flume::Sender;
7use libsqlite3_sys::*;
8use std::{
9    borrow::Cow,
10    ffi::{CStr, CString, c_char, c_int},
11    mem, ptr,
12    str::FromStr,
13    sync::{
14        Arc,
15        atomic::{AtomicPtr, Ordering},
16    },
17};
18use tank_core::{
19    AsQuery, Connection, Error, ErrorContext, Executor, Prepared, Query, QueryResult, RawQuery,
20    Result, Row, RowsAffected, error_message_from_ptr, send_value, stream::Stream, truncate_long,
21};
22use tokio::task::spawn_blocking;
23use url::Url;
24
25/// Wrapper for a SQLite `sqlite3` connection pointer used by the SQLite driver.
26///
27/// Provides helpers to prepare/execute statements and stream results into `tank_core` result types.
28pub struct SQLiteConnection {
29    pub(crate) connection: CBox<*mut sqlite3>,
30    pub(crate) url: Url,
31}
32
33impl SQLiteConnection {
34    pub fn last_error(&self) -> String {
35        unsafe {
36            let errcode = sqlite3_errcode(*self.connection);
37            format!(
38                "Error ({errcode}): {}",
39                error_message_from_ptr(&sqlite3_errmsg(*self.connection)),
40            )
41        }
42    }
43
44    pub(crate) fn do_run_prepared(
45        connection: *mut sqlite3,
46        statement: *mut sqlite3_stmt,
47        tx: Sender<Result<QueryResult>>,
48    ) {
49        unsafe {
50            let count = sqlite3_column_count(statement);
51            let labels = match (0..count)
52                .map(|i| extract_name(statement, i))
53                .collect::<Result<Arc<[_]>>>()
54            {
55                Ok(labels) => labels,
56                Err(error) => {
57                    send_value!(tx, Err(error.into()));
58                    return;
59                }
60            };
61            loop {
62                match sqlite3_step(statement) {
63                    SQLITE_BUSY => {
64                        continue;
65                    }
66                    SQLITE_DONE => {
67                        if sqlite3_stmt_readonly(statement) == 0 {
68                            send_value!(
69                                tx,
70                                Ok(QueryResult::Affected(RowsAffected {
71                                    rows_affected: Some(sqlite3_changes64(connection) as _),
72                                    last_affected_id: Some(sqlite3_last_insert_rowid(connection)),
73                                }))
74                            );
75                        }
76                        break;
77                    }
78                    SQLITE_ROW => {
79                        let values = match (0..count)
80                            .map(|i| extract_value(statement, i))
81                            .collect::<Result<_>>()
82                        {
83                            Ok(value) => value,
84                            Err(error) => {
85                                send_value!(tx, Err(error));
86                                return;
87                            }
88                        };
89                        send_value!(
90                            tx,
91                            Ok(QueryResult::Row(Row {
92                                labels: labels.clone(),
93                                values: values,
94                            }))
95                        )
96                    }
97                    _ => {
98                        send_value!(
99                            tx,
100                            Err(Error::msg(
101                                error_message_from_ptr(&sqlite3_errmsg(sqlite3_db_handle(
102                                    statement,
103                                )))
104                                .to_string(),
105                            ))
106                        );
107                        break;
108                    }
109                }
110            }
111        }
112    }
113
114    pub(crate) fn do_run_unprepared(
115        connection: *mut sqlite3,
116        sql: &str,
117        tx: Sender<Result<QueryResult>>,
118    ) {
119        unsafe {
120            let sql = sql.trim();
121            let mut it = sql.as_ptr() as *const c_char;
122            let mut len = sql.len();
123            loop {
124                let (statement, tail) = {
125                    let mut statement = SQLitePrepared::new(CBox::new(ptr::null_mut(), |p| {
126                        let rc = sqlite3_finalize(p);
127                        if rc != SQLITE_OK {
128                            return;
129                        }
130                    }));
131                    let mut sql_tail = ptr::null();
132                    let rc = sqlite3_prepare_v2(
133                        connection,
134                        it,
135                        len as c_int,
136                        &mut *statement.statement,
137                        &mut sql_tail,
138                    );
139                    if rc != SQLITE_OK {
140                        send_value!(
141                            tx,
142                            Err(Error::msg(
143                                error_message_from_ptr(&sqlite3_errmsg(connection)).to_string(),
144                            ))
145                        );
146                        return;
147                    }
148                    (statement, sql_tail)
149                };
150                Self::do_run_prepared(connection, statement.statement(), tx.clone());
151                len = if tail != ptr::null() {
152                    len - tail.offset_from_unsigned(it)
153                } else {
154                    0
155                };
156                if len == 0 {
157                    break;
158                }
159                it = tail;
160            }
161        };
162    }
163}
164
165impl Executor for SQLiteConnection {
166    type Driver = SQLiteDriver;
167
168    async fn do_prepare(&mut self, sql: String) -> Result<Query<SQLiteDriver>> {
169        let connection = AtomicPtr::new(*self.connection);
170        let context = format!("While preparing the query:\n{}", truncate_long!(sql));
171        let prepared = spawn_blocking(move || unsafe {
172            let connection = connection.load(Ordering::Relaxed);
173            let len = sql.len();
174            let sql = CString::new(sql.into_bytes())?;
175            let mut statement = CBox::new(ptr::null_mut(), |p| {
176                let db = sqlite3_db_handle(p);
177                let rc = sqlite3_finalize(p);
178                if rc != SQLITE_OK {
179                    let error = Error::msg(error_message_from_ptr(&sqlite3_errmsg(db)).to_string())
180                        .context("While finalizing a prepared statement");
181                    log::error!("{error:#}");
182                }
183            });
184            let mut tail = ptr::null();
185            let rc = sqlite3_prepare_v2(
186                connection,
187                sql.as_ptr(),
188                len as c_int,
189                &mut *statement,
190                &mut tail,
191            );
192            if rc != SQLITE_OK {
193                let error =
194                    Error::msg(error_message_from_ptr(&sqlite3_errmsg(connection)).to_string())
195                        .context(context);
196                log::error!("{:#}", error);
197                return Err(error);
198            }
199            if tail != ptr::null() && *tail != '\0' as i8 {
200                let error = Error::msg(format!(
201                    "Cannot prepare more than one statement at a time (remaining: {})",
202                    CStr::from_ptr(tail).to_str().unwrap_or("unprintable")
203                ))
204                .context(context);
205                log::error!("{:#}", error);
206                return Err(error);
207            }
208            Ok(statement)
209        })
210        .await?;
211        Ok(SQLitePrepared::new(prepared?).into())
212    }
213
214    fn run<'s>(
215        &'s mut self,
216        query: impl AsQuery<SQLiteDriver> + 's,
217    ) -> impl Stream<Item = Result<QueryResult>> + Send {
218        let mut query = query.as_query();
219        let context = Arc::new(format!("While running the query:\n{}", query.as_mut()));
220        let (tx, rx) = flume::unbounded::<Result<QueryResult>>();
221        let connection = AtomicPtr::new(*self.connection);
222        let mut owned = mem::take(query.as_mut());
223        let join = spawn_blocking(move || {
224            match &mut owned {
225                Query::Raw(RawQuery(sql)) => {
226                    Self::do_run_unprepared(connection.load(Ordering::Relaxed), sql, tx);
227                }
228                Query::Prepared(prepared) => {
229                    Self::do_run_prepared(
230                        connection.load(Ordering::Relaxed),
231                        prepared.statement(),
232                        tx,
233                    );
234                    let _ = prepared.clear_bindings();
235                }
236            }
237            owned
238        });
239        try_stream! {
240            while let Ok(result) = rx.recv_async().await {
241                yield result.map_err(|e| {
242                    let error = e.context(context.clone());
243                    log::error!("{:#}", error);
244                    error
245                })?;
246            }
247            *query.as_mut() = mem::take(&mut join.await?);
248        }
249    }
250}
251
252impl Connection for SQLiteConnection {
253    async fn connect(driver: &SQLiteDriver, url: Cow<'static, str>) -> Result<Self> {
254        let context = "While trying to connect to SQLite";
255        let url = Self::sanitize_url(driver, url)?;
256        let connection_string =
257            CString::from_str(&url.as_str().replacen("sqlite://", "file:", 1)).context(context)?;
258        let mut connection;
259        unsafe {
260            connection = CBox::new(ptr::null_mut(), |p| {
261                if sqlite3_close(p) != SQLITE_OK {
262                    let error = Error::msg(error_message_from_ptr(&sqlite3_errmsg(p)).to_string())
263                        .context("While closing the sqlite connection");
264                    log::error!("{error:#}");
265                }
266            });
267            let rc = sqlite3_open_v2(
268                connection_string.as_ptr(),
269                &mut *connection,
270                SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_URI,
271                ptr::null(),
272            );
273            if rc != SQLITE_OK {
274                let error =
275                    Error::msg(error_message_from_ptr(&sqlite3_errmsg(*connection)).to_string())
276                        .context(context);
277                log::error!("{:#}", error);
278                return Err(error);
279            }
280        }
281        Ok(Self { connection, url })
282    }
283
284    fn begin(&mut self) -> impl Future<Output = Result<SQLiteTransaction<'_>>> {
285        SQLiteTransaction::new(self)
286    }
287
288    async fn duplicate(&self) -> Result<SQLiteConnection>
289    where
290        Self: Sized,
291    {
292        Self::connect(&self.driver(), self.url.to_string().into()).await
293    }
294}
tank_sqlite/connection.rs

tank_sqlite/
connection.rs
