engine/src/script/db_bridge.rs - main - mcmare/RustAPI

engine/src/script/db_bridge.rs
268 строк10 KB
mcmare
Add Rhai scripting layer: host functions (db/hash/jwt), per-script transactions
18 июл 2026, 16:32
18 июл 2026, 16:324bfc229
use rhai::{Dynamic, Engine, EvalAltResult};
use serde_json::Value;
use std::{collections::HashMap, sync::Arc};

use super::convert::{dynamic_to_json, json_to_dynamic};
use crate::db::DbPool;

fn script_err(msg: impl Into<String>) -> Box<EvalAltResult> {
    Box::new(EvalAltResult::ErrorRuntime(Dynamic::from(msg.into()), rhai::Position::NONE))
}

/// A script author writes portable `?` placeholders regardless of the target driver;
/// Postgres needs `$1, $2, ...` instead, so it gets rewritten here. Naive but sufficient
/// for the simple parameterized queries scripts are expected to write: doesn't special-case
/// `?` that appears inside a quoted string literal other than skipping rewriting there.
fn rewrite_placeholders_for_postgres(sql: &str) -> String {
    let mut out = String::with_capacity(sql.len());
    let mut in_string = false;
    let mut n = 0u32;
    for c in sql.chars() {
        match c {
            '\'' => {
                in_string = !in_string;
                out.push(c);
            }
            '?' if !in_string => {
                n += 1;
                out.push('$');
                out.push_str(&n.to_string());
            }
            _ => out.push(c),
        }
    }
    out
}

fn identity(sql: &str) -> String {
    sql.to_string()
}

pub enum DbTx {
    Postgres(sqlx::Transaction<'static, sqlx::Postgres>),
    MySql(sqlx::Transaction<'static, sqlx::MySql>),
    Sqlite(sqlx::Transaction<'static, sqlx::Sqlite>),
}

/// Per-script-execution DB state: lazily-begun transactions, one per connection touched,
/// committed together on success or rolled back together on failure by `finish`.
#[derive(Clone)]
pub struct DbContext {
    pools: Arc<HashMap<String, DbPool>>,
    tx: Arc<std::sync::Mutex<HashMap<String, DbTx>>>,
}

impl DbContext {
    pub fn new(pools: Arc<HashMap<String, DbPool>>) -> Self {
        Self {
            pools,
            tx: Arc::new(std::sync::Mutex::new(HashMap::new())),
        }
    }

    pub async fn finish(&self, commit: bool) -> anyhow::Result<()> {
        let txs: Vec<(String, DbTx)> = {
            let mut guard = self.tx.lock().unwrap();
            guard.drain().collect()
        };
        for (_, tx) in txs {
            match tx {
                DbTx::Postgres(t) => {
                    if commit {
                        t.commit().await?
                    } else {
                        t.rollback().await?
                    }
                }
                DbTx::MySql(t) => {
                    if commit {
                        t.commit().await?
                    } else {
                        t.rollback().await?
                    }
                }
                DbTx::Sqlite(t) => {
                    if commit {
                        t.commit().await?
                    } else {
                        t.rollback().await?
                    }
                }
            }
        }
        Ok(())
    }

    fn ensure_tx(&self, connection: &str) -> Result<(), Box<EvalAltResult>> {
        let mut guard = self.tx.lock().unwrap();
        if guard.contains_key(connection) {
            return Ok(());
        }
        let pool = self
            .pools
            .get(connection)
            .ok_or_else(|| script_err(format!("unknown connection '{connection}'")))?;
        let handle = tokio::runtime::Handle::current();
        let tx = handle
            .block_on(async {
                match pool {
                    DbPool::Postgres(p) => p.begin().await.map(DbTx::Postgres),
                    DbPool::MySql(p) => p.begin().await.map(DbTx::MySql),
                    DbPool::Sqlite(p) => p.begin().await.map(DbTx::Sqlite),
                }
            })
            .map_err(|e| script_err(format!("failed to begin transaction on '{connection}': {e}")))?;
        guard.insert(connection.to_string(), tx);
        Ok(())
    }

    pub fn query_rows(&self, connection: &str, sql: &str, params: &[Dynamic]) -> Result<Vec<Value>, Box<EvalAltResult>> {
        self.ensure_tx(connection)?;
        let mut guard = self.tx.lock().unwrap();
        let tx = guard.get_mut(connection).expect("just ensured");
        let handle = tokio::runtime::Handle::current();
        let result = match tx {
            DbTx::Postgres(t) => handle.block_on(pg::run_fetch_all(&mut **t, sql, params)),
            DbTx::MySql(t) => handle.block_on(mysql::run_fetch_all(&mut **t, sql, params)),
            DbTx::Sqlite(t) => handle.block_on(sqlite::run_fetch_all(&mut **t, sql, params)),
        };
        result.map_err(script_err)
    }

    pub fn execute(&self, connection: &str, sql: &str, params: &[Dynamic]) -> Result<u64, Box<EvalAltResult>> {
        self.ensure_tx(connection)?;
        let mut guard = self.tx.lock().unwrap();
        let tx = guard.get_mut(connection).expect("just ensured");
        let handle = tokio::runtime::Handle::current();
        let result = match tx {
            DbTx::Postgres(t) => handle.block_on(pg::run_execute(&mut **t, sql, params)),
            DbTx::MySql(t) => handle.block_on(mysql::run_execute(&mut **t, sql, params)),
            DbTx::Sqlite(t) => handle.block_on(sqlite::run_execute(&mut **t, sql, params)),
        };
        result.map_err(script_err)
    }
}

macro_rules! impl_backend {
    ($mod_name:ident, $db:ty, $row:ty, $prepare:expr) => {
        mod $mod_name {
            use super::*;
            use sqlx::{Column, Row};

            fn decode_cell(row: &$row, idx: usize) -> Value {
                macro_rules! attempt {
                    ($t:ty, $conv:expr) => {
                        match row.try_get::<Option<$t>, usize>(idx) {
                            Ok(Some(v)) => return $conv(v),
                            Ok(None) => return Value::Null,
                            Err(_) => {}
                        }
                    };
                }
                attempt!(i64, |v: i64| Value::Number(v.into()));
                attempt!(f64, |v: f64| serde_json::Number::from_f64(v)
                    .map(Value::Number)
                    .unwrap_or(Value::Null));
                attempt!(bool, Value::Bool);
                attempt!(uuid::Uuid, |v: uuid::Uuid| Value::String(v.to_string()));
                attempt!(chrono::DateTime<chrono::Utc>, |v: chrono::DateTime<chrono::Utc>| Value::String(
                    v.to_rfc3339()
                ));
                attempt!(String, Value::String);
                attempt!(Value, |v: Value| v);
                Value::Null
            }

            fn row_to_json(row: &$row) -> Value {
                let mut obj = serde_json::Map::new();
                for col in row.columns() {
                    obj.insert(col.name().to_string(), decode_cell(row, col.ordinal()));
                }
                Value::Object(obj)
            }

            pub(super) async fn run_fetch_all<'e, E>(exec: E, sql: &str, params: &[Dynamic]) -> Result<Vec<Value>, String>
            where
                E: sqlx::Executor<'e, Database = $db>,
            {
                let prepared = $prepare(sql);
                let mut q = sqlx::query(&prepared);
                for p in params {
                    q = if let Some(i) = p.clone().try_cast::<i64>() {
                        q.bind(i)
                    } else if let Some(f) = p.clone().try_cast::<f64>() {
                        q.bind(f)
                    } else if let Some(b) = p.clone().try_cast::<bool>() {
                        q.bind(b)
                    } else if let Some(s) = p.clone().try_cast::<rhai::ImmutableString>() {
                        q.bind(s.to_string())
                    } else if p.is_unit() {
                        q.bind(Option::<String>::None)
                    } else {
                        q.bind(dynamic_to_json(p))
                    };
                }
                let rows = q.fetch_all(exec).await.map_err(|e| e.to_string())?;
                Ok(rows.iter().map(row_to_json).collect())
            }

            pub(super) async fn run_execute<'e, E>(exec: E, sql: &str, params: &[Dynamic]) -> Result<u64, String>
            where
                E: sqlx::Executor<'e, Database = $db>,
            {
                let prepared = $prepare(sql);
                let mut q = sqlx::query(&prepared);
                for p in params {
                    q = if let Some(i) = p.clone().try_cast::<i64>() {
                        q.bind(i)
                    } else if let Some(f) = p.clone().try_cast::<f64>() {
                        q.bind(f)
                    } else if let Some(b) = p.clone().try_cast::<bool>() {
                        q.bind(b)
                    } else if let Some(s) = p.clone().try_cast::<rhai::ImmutableString>() {
                        q.bind(s.to_string())
                    } else if p.is_unit() {
                        q.bind(Option::<String>::None)
                    } else {
                        q.bind(dynamic_to_json(p))
                    };
                }
                let result = q.execute(exec).await.map_err(|e| e.to_string())?;
                Ok(result.rows_affected())
            }
        }
    };
}

impl_backend!(pg, sqlx::Postgres, sqlx::postgres::PgRow, rewrite_placeholders_for_postgres);
impl_backend!(mysql, sqlx::MySql, sqlx::mysql::MySqlRow, identity);
impl_backend!(sqlite, sqlx::Sqlite, sqlx::sqlite::SqliteRow, identity);

/// The `db` object exposed to scripts: `db.query(...)`, `db.query_one(...)`, `db.execute(...)`.
#[derive(Clone)]
pub struct DbApi {
    pub ctx: DbContext,
}

impl DbApi {
    fn query(&mut self, connection: &str, sql: &str, params: rhai::Array) -> Result<rhai::Array, Box<EvalAltResult>> {
        let rows = self.ctx.query_rows(connection, sql, &params)?;
        Ok(rows.iter().map(json_to_dynamic).collect())
    }

    fn query_one(&mut self, connection: &str, sql: &str, params: rhai::Array) -> Result<Dynamic, Box<EvalAltResult>> {
        let rows = self.ctx.query_rows(connection, sql, &params)?;
        Ok(rows.into_iter().next().map(|v| json_to_dynamic(&v)).unwrap_or(Dynamic::UNIT))
    }

    fn execute(&mut self, connection: &str, sql: &str, params: rhai::Array) -> Result<i64, Box<EvalAltResult>> {
        Ok(self.ctx.execute(connection, sql, &params)? as i64)
    }
}

pub fn register(engine: &mut Engine) {
    engine.register_type_with_name::<DbApi>("DbApi");
    engine.register_fn("query", DbApi::query);
    engine.register_fn("query_one", DbApi::query_one);
    engine.register_fn("execute", DbApi::execute);
}