spin_sqlite_libsql/

lib.rs

use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use spin_factor_sqlite::{Connection, QueryAsyncResult};
use spin_world::spin::sqlite3_1_0::sqlite as v3;
use spin_world::spin::sqlite3_1_0::sqlite::{self, RowResult};
use tokio::sync::OnceCell;

/// A lazy wrapper around a [`LibSqlConnection`] that implements the [`Connection`] trait.
pub struct LazyLibSqlConnection {
    url: String,
    token: String,
    // Since the libSQL client can only be created asynchronously, we wait until
    // we're in the `Connection` implementation to create. Since we only want to do
    // this once, we use a `OnceCell` to store it.
    inner: OnceCell<Arc<LibSqlConnection>>,
}

impl LazyLibSqlConnection {
    pub fn new(url: String, token: String) -> Self {
        Self {
            url,
            token,
            inner: OnceCell::new(),
        }
    }

    pub async fn get_or_create_connection(&self) -> Result<&Arc<LibSqlConnection>, v3::Error> {
        self.inner
            .get_or_try_init(|| async {
                LibSqlConnection::create(self.url.clone(), self.token.clone())
                    .await
                    .context("failed to create SQLite client")
                    .map(Arc::new)
            })
            .await
            .map_err(|_| v3::Error::InvalidConnection)
    }
}

#[async_trait]
impl Connection for LazyLibSqlConnection {
    async fn query(
        &self,
        query: &str,
        parameters: Vec<v3::Value>,
        max_result_bytes: usize,
    ) -> Result<v3::QueryResult, v3::Error> {
        let client = self.get_or_create_connection().await?;
        client.query(query, parameters, max_result_bytes).await
    }

    async fn query_async(
        &self,
        query: &str,
        parameters: Vec<v3::Value>,
        max_result_bytes: usize,
    ) -> Result<QueryAsyncResult, v3::Error> {
        let client = self.get_or_create_connection().await?.clone();
        let query = query.to_string();

        let (cols_tx, cols_rx) = tokio::sync::oneshot::channel();
        let (rows_tx, rows_rx) = tokio::sync::mpsc::channel(4);
        let (err_tx, err_rx) = tokio::sync::oneshot::channel();

        let the_work = async move {
            let result = client
                .inner
                .query(&query, convert_parameters(&parameters))
                .await
                .map_err(|e| v3::Error::Io(e.to_string()));

            let mut rows = match result {
                Ok(r) => r,
                Err(e) => {
                    _ = cols_tx.send(Default::default());
                    return Err(e);
                }
            };

            let columns = columns(&rows);
            cols_tx
                .send(columns)
                .map_err(|_| v3::Error::Io("column send error".into()))?;

            let column_count = rows.column_count();

            loop {
                let row = match rows.next().await.map_err(io_error_v3)? {
                    Some(r) => r,
                    None => break,
                };
                let row = convert_row(row, column_count);
                if row.values.iter().map(|v| v.memory_size()).sum::<usize>() > max_result_bytes {
                    return Err(v3::Error::Io(format!(
                        "query result exceeds limit of {max_result_bytes} bytes"
                    )));
                }
                rows_tx
                    .send(row)
                    .await
                    .map_err(|_| v3::Error::Io("column send error".into()))?;
            }

            Ok(())
        };
        tokio::spawn(async move {
            let res = the_work.await;
            _ = err_tx.send(res);
        });

        let columns = cols_rx.await.map_err(|e| v3::Error::Io(e.to_string()))?;

        Ok(QueryAsyncResult {
            columns,
            rows: rows_rx,
            error: err_rx,
        })
    }

    async fn execute_batch(&self, statements: &str) -> anyhow::Result<()> {
        let client = self.get_or_create_connection().await?;
        client.execute_batch(statements).await
    }

    async fn changes(&self) -> Result<u64, sqlite::Error> {
        let client = self.get_or_create_connection().await?;
        Ok(client.changes())
    }

    async fn last_insert_rowid(&self) -> Result<i64, sqlite::Error> {
        let client = self.get_or_create_connection().await?;
        Ok(client.last_insert_rowid())
    }

    fn summary(&self) -> Option<String> {
        Some(format!("libSQL at {}", self.url))
    }
}

/// An open connection to a libSQL server.
#[derive(Clone)]
pub struct LibSqlConnection {
    inner: libsql::Connection,
}

impl LibSqlConnection {
    pub async fn create(url: String, token: String) -> anyhow::Result<Self> {
        let db = libsql::Builder::new_remote(url, token).build().await?;
        let inner = db.connect()?;
        Ok(Self { inner })
    }
}

impl LibSqlConnection {
    pub async fn query(
        &self,
        query: &str,
        parameters: Vec<sqlite::Value>,
        max_result_bytes: usize,
    ) -> Result<sqlite::QueryResult, sqlite::Error> {
        let result = self
            .inner
            .query(query, convert_parameters(&parameters))
            .await
            .map_err(|e| sqlite::Error::Io(e.to_string()))?;

        Ok(sqlite::QueryResult {
            columns: columns(&result),
            rows: convert_rows(result, max_result_bytes)
                .await
                .map_err(|e| sqlite::Error::Io(e.to_string()))?,
        })
    }

    pub async fn execute_batch(&self, statements: &str) -> anyhow::Result<()> {
        self.inner.execute_batch(statements).await?;

        Ok(())
    }

    pub fn changes(&self) -> u64 {
        self.inner.changes()
    }

    pub fn last_insert_rowid(&self) -> i64 {
        self.inner.last_insert_rowid()
    }
}

fn columns(rows: &libsql::Rows) -> Vec<String> {
    (0..rows.column_count())
        .map(|index| rows.column_name(index).unwrap_or("").to_owned())
        .collect()
}

async fn convert_rows(
    mut rows: libsql::Rows,
    max_result_bytes: usize,
) -> anyhow::Result<Vec<RowResult>> {
    let mut result_rows = vec![];

    let column_count = rows.column_count();
    let mut byte_count = 0;
    while let Some(row) = rows.next().await? {
        let row = convert_row(row, column_count);
        byte_count += row.values.iter().map(|v| v.memory_size()).sum::<usize>();
        if byte_count > max_result_bytes {
            anyhow::bail!("query result exceeds limit of {max_result_bytes} bytes")
        }
        result_rows.push(row);
    }

    Ok(result_rows)
}

fn convert_row(row: libsql::Row, column_count: i32) -> RowResult {
    let values = (0..column_count)
        .map(|index| convert_value(row.get_value(index).unwrap()))
        .collect();
    RowResult { values }
}

fn convert_value(v: libsql::Value) -> sqlite::Value {
    use libsql::Value;

    match v {
        Value::Null => sqlite::Value::Null,
        Value::Integer(value) => sqlite::Value::Integer(value),
        Value::Real(value) => sqlite::Value::Real(value),
        Value::Text(value) => sqlite::Value::Text(value),
        Value::Blob(value) => sqlite::Value::Blob(value),
    }
}

fn convert_parameters(parameters: &[sqlite::Value]) -> Vec<libsql::Value> {
    use libsql::Value;

    parameters
        .iter()
        .map(|v| match v {
            sqlite::Value::Integer(value) => Value::Integer(*value),
            sqlite::Value::Real(value) => Value::Real(*value),
            sqlite::Value::Text(t) => Value::Text(t.clone()),
            sqlite::Value::Blob(b) => Value::Blob(b.clone()),
            sqlite::Value::Null => Value::Null,
        })
        .collect()
}

fn io_error_v3(err: libsql::Error) -> v3::Error {
    v3::Error::Io(err.to_string())
}