nautilus-orm-connector 0.1.7

//! SQLite executor implementation.

use crate::error::{ConnectorError as Error, Result};
use crate::{Executor, Row, SqliteRowStream};
use nautilus_core::Value;
use nautilus_dialect::Sql;
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};

/// SQLite executor using sqlx.
///
/// Manages a connection pool and executes queries against SQLite databases.
///
/// ## Example
///
/// ```rust,ignore
/// use nautilus_connector::SqliteExecutor;
///
/// #[tokio::main]
/// async fn main() -> nautilus_core::Result<()> {
///     // File-based database
///     let executor = SqliteExecutor::new("sqlite:mydb.db").await?;
///     // Or in-memory database
///     let executor = SqliteExecutor::new("sqlite::memory:").await?;
///     // Use executor to run queries...
///     Ok(())
/// }
/// ```
pub struct SqliteExecutor {
    pool: SqlitePool,
}

impl SqliteExecutor {
    /// Create a new SQLite executor with a connection pool.
    ///
    /// ## Parameters
    ///
    /// - `url`: SQLite connection URL (e.g., `sqlite:mydb.db` or `sqlite::memory:`)
    ///
    /// ## Errors
    ///
    /// Returns `ConnectorError::Connection` if the pool cannot be created.
    pub async fn new(url: &str) -> Result<Self> {
        let pool = SqlitePoolOptions::new()
            .max_connections(5)
            .connect(url)
            .await
            .map_err(|e| Error::connection(e, "Failed to connect to SQLite"))?;

        Ok(Self { pool })
    }

    /// Get a reference to the underlying connection pool.
    pub fn pool(&self) -> &SqlitePool {
        &self.pool
    }

    /// Execute a raw SQL statement with no result rows (e.g., DDL).
    pub async fn execute_raw(&self, sql: &str) -> Result<()> {
        sqlx::query(sql)
            .execute(&self.pool)
            .await
            .map(|_| ())
            .map_err(|e| Error::database(e, "DDL error"))
    }

    impl_execute_affected!();
}

/// [`Executor`] implementation backed by a SQLite connection pool.
impl Executor for SqliteExecutor {
    type Row<'conn>
        = Row
    where
        Self: 'conn;
    type RowStream<'conn>
        = SqliteRowStream
    where
        Self: 'conn;

    fn execute<'conn>(&'conn self, sql: &'conn Sql) -> Self::RowStream<'conn> {
        let pool = self.pool.clone();
        let sql_text = sql.text.clone();
        let params = sql.params.clone();

        let stream = async_stream::stream! {
            let mut conn = match pool.acquire().await {
                Ok(c) => c,
                Err(e) => {
                    yield Err(Error::connection(e, "Failed to acquire connection"));
                    return;
                }
            };

            let mut query = sqlx::query(&sql_text);
            for param in &params {
                query = match bind_value(query, param) {
                    Ok(q) => q,
                    Err(e) => {
                        yield Err(e);
                        return;
                    }
                };
            }

            let sqlite_rows = match query.fetch_all(&mut *conn).await {
                Ok(rows) => rows,
                Err(e) => {
                    yield Err(Error::database(e, "Query execution failed"));
                    return;
                }
            };

            drop(conn);

            for sqlite_row in sqlite_rows {
                match crate::sqlite_stream::decode_row_internal(sqlite_row) {
                    Ok(row) => yield Ok(row),
                    Err(e) => yield Err(e),
                }
            }
        };

        SqliteRowStream::new_from_stream(Box::pin(stream))
    }

    fn execute_and_fetch<'conn>(
        &'conn self,
        mutation: &'conn Sql,
        fetch: &'conn Sql,
    ) -> Self::RowStream<'conn> {
        let pool = self.pool.clone();
        let mutation_text = mutation.text.clone();
        let mutation_params = mutation.params.clone();
        let fetch_text = fetch.text.clone();
        let fetch_params = fetch.params.clone();

        let stream = async_stream::stream! {
            use sqlx::Executor as _;

            let mut conn = match pool.acquire().await {
                Ok(c) => c,
                Err(e) => {
                    yield Err(Error::connection(e, "Failed to acquire connection"));
                    return;
                }
            };

            let mut mutation_query = sqlx::query(&mutation_text);
            for param in &mutation_params {
                mutation_query = match bind_value(mutation_query, param) {
                    Ok(q) => q,
                    Err(e) => {
                        yield Err(e);
                        return;
                    }
                };
            }

            if let Err(e) = (&mut *conn).execute(mutation_query).await {
                yield Err(Error::database(e, "Mutation failed"));
                return;
            }

            let mut fetch_query = sqlx::query(&fetch_text);
            for param in &fetch_params {
                fetch_query = match bind_value(fetch_query, param) {
                    Ok(q) => q,
                    Err(e) => {
                        yield Err(e);
                        return;
                    }
                };
            }

            let sqlite_rows = match fetch_query.fetch_all(&mut *conn).await {
                Ok(rows) => rows,
                Err(e) => {
                    yield Err(Error::database(e, "Fetch failed"));
                    return;
                }
            };

            drop(conn);

            for sqlite_row in sqlite_rows {
                match crate::sqlite_stream::decode_row_internal(sqlite_row) {
                    Ok(row) => yield Ok(row),
                    Err(e) => yield Err(e),
                }
            }
        };

        SqliteRowStream::new_from_stream(Box::pin(stream))
    }
}

/// Binds a [`Value`] to a SQLite sqlx query as a typed parameter.
///
/// `Decimal`, `DateTime`, and `Uuid` are serialized to strings because SQLite
/// has no native support for those types.  Arrays are serialized as JSON strings.
/// Note: `SqliteArguments` carries a lifetime parameter `'q`, unlike the PG/MySQL
/// argument types, which is why this function cannot be made generic across all
/// three backends with a single signature.
pub(crate) fn bind_value<'q>(
    query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>,
    value: &'q Value,
) -> Result<sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>> {
    match value {
        Value::Null => Ok(query.bind(None::<String>)),
        Value::Bool(b) => Ok(query.bind(b)),
        Value::I32(i) => Ok(query.bind(i)),
        Value::I64(i) => Ok(query.bind(i)),
        Value::F64(f) => Ok(query.bind(f)),
        Value::Decimal(d) => Ok(query.bind(d.to_string())),
        Value::DateTime(dt) => Ok(query.bind(dt.format("%Y-%m-%dT%H:%M:%S%.f").to_string())),
        Value::Uuid(u) => Ok(query.bind(u.to_string())),
        Value::String(s) => Ok(query.bind(s.as_str())),
        Value::Bytes(b) => Ok(query.bind(b.as_slice())),
        Value::Json(j) => Ok(query.bind(j.to_string())),
        Value::Array(_) => Ok(query.bind(crate::utils::value_to_json(value).to_string())),
        Value::Enum { value, .. } => Ok(query.bind(value.as_str())),
        Value::Array2D(_) => Ok(query.bind(crate::utils::value_to_json(value).to_string())),
    }
}