use async_trait::async_trait;
use sqlx::sqlite::{SqlitePool, SqliteRow};
use sqlx::{Transaction, Row, Column};
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::core::{RiResult, RiError};
use crate::database::{
RiDatabase, RiDatabaseConfig, DatabaseType,
RiDBResult, RiDBRow
};
#[derive(Clone)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
#[allow(dead_code)]
pub struct SQLiteDatabase {
pool: SqlitePool,
config: RiDatabaseConfig,
}
impl SQLiteDatabase {
pub async fn new(database_url: &str, config: RiDatabaseConfig) -> Result<Self, RiError> {
let pool = SqlitePool::connect(database_url)
.await
.map_err(|e| RiError::Other(format!("Failed to connect to SQLite database: {}", e)))?;
Ok(Self { pool, config })
}
fn row_to_ri_row(row: &SqliteRow) -> RiDBRow {
let columns: Vec<String> = (0..row.len())
.map(|i| row.column(i).name().to_string())
.collect();
let values: Vec<Option<serde_json::Value>> = (0..row.len())
.map(|idx| Self::value_to_json(row, idx))
.collect();
RiDBRow { columns, values }
}
fn value_to_json(row: &SqliteRow, idx: usize) -> Option<serde_json::Value> {
match row.try_get::<i64, _>(idx) {
Ok(v) => Some(serde_json::json!(v)),
Err(_) => {
match row.try_get::<f64, _>(idx) {
Ok(v) => Some(serde_json::json!(v)),
Err(_) => {
match row.try_get::<String, _>(idx) {
Ok(v) => Some(serde_json::json!(v)),
Err(_) => {
match row.try_get::<bool, _>(idx) {
Ok(v) => Some(serde_json::json!(v)),
Err(_) => {
match row.try_get::<Vec<u8>, _>(idx) {
Ok(v) => Some(serde_json::json!(v)),
Err(_) => None,
}
}
}
}
}
}
}
}
}
}
}
#[async_trait]
impl RiDatabase for SQLiteDatabase {
fn database_type(&self) -> DatabaseType {
DatabaseType::SQLite
}
async fn execute(&self, sql: &str) -> RiResult<u64> {
let result = sqlx::query::<sqlx::Sqlite>(sql)
.execute(&self.pool)
.await
.map_err(|e| RiError::Other(format!("SQLite execute error: {}", e)))?;
Ok(result.rows_affected())
}
async fn query(&self, sql: &str) -> RiResult<RiDBResult> {
let rows = sqlx::query::<sqlx::Sqlite>(sql)
.fetch_all(&self.pool)
.await
.map_err(|e| RiError::Other(format!("SQLite query error: {}", e)))?;
let ri_rows: Vec<RiDBRow> = rows.iter()
.map(|row| Self::row_to_ri_row(row))
.collect();
Ok(RiDBResult::with_rows(ri_rows))
}
async fn query_one(&self, sql: &str) -> RiResult<Option<RiDBRow>> {
let row = sqlx::query::<sqlx::Sqlite>(sql)
.fetch_optional(&self.pool)
.await
.map_err(|e| RiError::Other(format!("SQLite query_one error: {}", e)))?;
Ok(row.map(|r| Self::row_to_ri_row(&r)))
}
async fn ping(&self) -> RiResult<bool> {
sqlx::query::<sqlx::Sqlite>("SELECT 1")
.fetch_one(&self.pool)
.await
.map(|_| true)
.map_err(|e| RiError::Other(format!("SQLite ping error: {}", e)))
}
fn is_connected(&self) -> bool {
!self.pool.is_closed()
}
async fn close(&self) -> RiResult<()> {
self.pool.close().await;
Ok(())
}
async fn batch_execute(&self, sql: &str, params: &[Vec<serde_json::Value>]) -> RiResult<Vec<u64>> {
let mut results = Vec::with_capacity(params.len());
for param_set in params {
let result = self.execute_with_params(sql, param_set).await?;
results.push(result);
}
Ok(results)
}
async fn batch_query(&self, sql: &str, params: &[Vec<serde_json::Value>]) -> RiResult<Vec<RiDBResult>> {
let mut results = Vec::with_capacity(params.len());
for param_set in params {
let result = self.query_with_params(sql, param_set).await?;
results.push(result);
}
Ok(results)
}
async fn execute_with_params(&self, sql: &str, params: &[serde_json::Value]) -> RiResult<u64> {
let mut query = sqlx::query::<sqlx::Sqlite>(sql);
for param in params {
let param_str = param.to_string();
query = query.bind(param_str);
}
let result = query
.execute(&self.pool)
.await
.map_err(|e| RiError::Other(format!("SQLite execute_with_params error: {}", e)))?;
Ok(result.rows_affected())
}
async fn query_with_params(&self, sql: &str, params: &[serde_json::Value]) -> RiResult<RiDBResult> {
let mut query = sqlx::query::<sqlx::Sqlite>(sql);
for param in params {
let param_str = param.to_string();
query = query.bind(param_str);
}
let rows = query
.fetch_all(&self.pool)
.await
.map_err(|e| RiError::Other(format!("SQLite query_with_params error: {}", e)))?;
let ri_rows: Vec<RiDBRow> = rows.iter()
.map(|row| Self::row_to_ri_row(row))
.collect();
Ok(RiDBResult::with_rows(ri_rows))
}
async fn transaction(&self) -> RiResult<Box<dyn crate::database::RiDatabaseTransaction>> {
let tx = self.pool.begin().await
.map_err(|e| RiError::Other(format!("SQLite transaction begin error: {}", e)))?;
Ok(Box::new(SQLiteTransaction::new(tx)))
}
}
struct SQLiteTransaction {
tx: Arc<Mutex<Option<Transaction<'static, sqlx::Sqlite>>>>,
}
impl SQLiteTransaction {
pub fn new(tx: Transaction<'static, sqlx::Sqlite>) -> Self {
Self {
tx: Arc::new(Mutex::new(Some(tx))),
}
}
}
#[async_trait::async_trait]
impl crate::database::RiDatabaseTransaction for SQLiteTransaction {
async fn execute(&self, sql: &str) -> RiResult<u64> {
let mut guard = self.tx.lock().await;
let tx = guard.as_mut()
.ok_or_else(|| RiError::Other("SQLite transaction already closed".to_string()))?;
let result = sqlx::query::<sqlx::Sqlite>(sql)
.execute(&mut **tx)
.await
.map_err(|e| RiError::Other(format!("SQLite transaction execute error: {}", e)))?;
Ok(result.rows_affected())
}
async fn query(&self, sql: &str) -> RiResult<RiDBResult> {
let mut guard = self.tx.lock().await;
let tx = guard.as_mut()
.ok_or_else(|| RiError::Other("SQLite transaction already closed".to_string()))?;
let rows = sqlx::query::<sqlx::Sqlite>(sql)
.fetch_all(&mut **tx)
.await
.map_err(|e| RiError::Other(format!("SQLite transaction query error: {}", e)))?;
let ri_rows: Vec<RiDBRow> = rows.iter()
.map(|row| SQLiteDatabase::row_to_ri_row(row))
.collect();
Ok(RiDBResult::with_rows(ri_rows))
}
async fn commit(&self) -> RiResult<()> {
let mut guard = self.tx.lock().await;
let tx = guard.take()
.ok_or_else(|| RiError::Other("SQLite transaction already closed".to_string()))?;
tx.commit().await
.map_err(|e| RiError::Other(format!("SQLite transaction commit error: {}", e)))
}
async fn rollback(&self) -> RiResult<()> {
let mut guard = self.tx.lock().await;
let tx = guard.take()
.ok_or_else(|| RiError::Other("SQLite transaction already closed".to_string()))?;
tx.rollback().await
.map_err(|e| RiError::Other(format!("SQLite transaction rollback error: {}", e)))
}
async fn close(&self) -> RiResult<()> {
self.rollback().await
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl SQLiteDatabase {
#[staticmethod]
pub fn from_path(path: &str, max_connections: u32) -> Result<Self, pyo3::PyErr> {
let rt = tokio::runtime::Runtime::new()
.map_err(|e| pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to create Tokio runtime: {}", e),
))?;
let url = format!("sqlite:{}", path);
rt.block_on(async {
let pool = SqlitePool::connect(&url)
.await
.map_err(|e| pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to connect to SQLite database: {}", e),
))?;
let db_config = RiDatabaseConfig::sqlite(path)
.max_connections(max_connections)
.build();
Ok(Self { pool, config: db_config })
})
}
pub fn execute_sync(&self, sql: &str) -> Result<u64, RiError> {
let rt = tokio::runtime::Runtime::new()
.map_err(|e| RiError::Other(format!("Failed to create Tokio runtime: {}", e)))?;
rt.block_on(async {
self.execute(sql).await
})
}
pub fn query_sync(&self, sql: &str) -> Result<RiDBResult, RiError> {
let rt = tokio::runtime::Runtime::new()
.map_err(|e| RiError::Other(format!("Failed to create Tokio runtime: {}", e)))?;
rt.block_on(async {
self.query(sql).await
})
}
pub fn ping_sync(&self) -> Result<bool, RiError> {
let rt = tokio::runtime::Runtime::new()
.map_err(|e| RiError::Other(format!("Failed to create Tokio runtime: {}", e)))?;
rt.block_on(async {
self.ping().await
})
}
}