use akar_extension::{Extension, ExtensionContext};
use std::sync::Arc;
#[cfg(feature = "bundled")]
fn sqlite_value_to_string(val: &rusqlite::types::Value) -> String {
match val {
rusqlite::types::Value::Null => "NULL".into(),
rusqlite::types::Value::Integer(i) => i.to_string(),
rusqlite::types::Value::Real(f) => f.to_string(),
rusqlite::types::Value::Text(s) => s.clone(),
rusqlite::types::Value::Blob(b) => format!("<blob:{} bytes>", b.len()),
}
}
pub struct SqliteExtension;
impl Default for SqliteExtension {
fn default() -> Self {
Self::new()
}
}
impl SqliteExtension {
pub fn new() -> Self {
Self
}
}
impl Extension for SqliteExtension {
fn name(&self) -> &'static str {
"SQLITE"
}
fn load(&self, context: &ExtensionContext) -> Result<(), String> {
#[allow(unused_imports)]
use akar_function::Value;
use akar_function::registry::{ScalarFunction, TableFunction};
#[cfg(feature = "bundled")]
{
let query_fn: Arc<dyn Fn(&[Value]) -> Result<Value, String> + Send + Sync> = Arc::new(|args| {
if args.len() < 2 {
return Err("sqlite_query requires (path, sql) arguments".into());
}
let path = match &args[0] {
Value::String(s) => s.clone(),
_ => return Err("sqlite_query: first argument must be a path string".into()),
};
let sql = match &args[1] {
Value::String(s) => s.clone(),
_ => return Err("sqlite_query: second argument must be a SQL string".into()),
};
let conn =
rusqlite::Connection::open(&path).map_err(|e| format!("Failed to open SQLite DB '{path}': {e}"))?;
let mut stmt = conn.prepare(&sql).map_err(|e| format!("SQLite prepare error: {e}"))?;
let col_count = stmt.column_count();
let mut rows = stmt.query([]).map_err(|e| format!("SQLite query error: {e}"))?;
let mut parts = Vec::new();
while let Some(row) = rows.next().map_err(|e| format!("SQLite row error: {e}"))? {
for i in 0..col_count {
let val: rusqlite::types::Value = row
.get::<_, rusqlite::types::Value>(i)
.unwrap_or(rusqlite::types::Value::Null);
parts.push(sqlite_value_to_string(&val));
}
}
if parts.is_empty() {
Ok(Value::String("(empty)".into()))
} else {
Ok(Value::String(parts.join(",")))
}
});
context.register_scalar_function(
"sqlite_query",
ScalarFunction::CustomScalar {
name: "sqlite_query".into(),
execute: query_fn,
},
);
let scan_fn: Arc<dyn Fn(&[Value], &mut akar_function::DataChunk) -> Result<(), String> + Send + Sync> =
Arc::new(|args, chunk| {
if args.len() < 2 {
return Err("sqlite_scan requires (path, table) arguments".into());
}
let path = match &args[0] {
Value::String(s) => s.clone(),
_ => return Err("sqlite_scan: first argument must be a path string".into()),
};
let table = match &args[1] {
Value::String(s) => s.clone(),
_ => return Err("sqlite_scan: second argument must be a table name".into()),
};
if chunk.size > 0 {
return Ok(());
}
let conn = rusqlite::Connection::open(&path)
.map_err(|e| format!("Failed to open SQLite DB '{path}': {e}"))?;
let sql = format!("SELECT * FROM {}", akar_common::extension_utils::quote_sql_table_name(&table));
let mut stmt = conn.prepare(&sql).map_err(|e| format!("SQLite prepare error: {e}"))?;
let col_count = stmt.column_count();
let names: Vec<String> = (0..col_count)
.map(|i| stmt.column_name(i).map(str::to_string).unwrap_or_default())
.collect();
let mut rows = stmt.query([]).map_err(|e| format!("SQLite query error: {e}"))?;
let mut columns: Vec<Vec<Option<String>>> = vec![Vec::new(); col_count];
while let Some(row) = rows.next().map_err(|e| format!("SQLite row error: {e}"))? {
for i in 0..col_count {
let val: rusqlite::types::Value = row
.get::<_, rusqlite::types::Value>(i)
.unwrap_or(rusqlite::types::Value::Null);
columns[i].push(Some(sqlite_value_to_string(&val)));
}
}
chunk.fields.clear();
chunk.field_types.clear();
chunk.field_names.clear();
for (col, name) in columns.into_iter().zip(names) {
chunk.fields.push(
std::sync::Arc::new(arrow::array::StringArray::from_iter(
col.iter().map(|o| o.as_deref()),
)) as arrow::array::ArrayRef,
);
chunk.field_types.push(akar_common::types::PhysicalTypeID::String);
chunk.field_names.push(name);
}
chunk.size = chunk.fields.first().map(|f| f.len()).unwrap_or(0);
Ok(())
});
context.register_table_function(
"sqlite_scan",
TableFunction::CustomTable {
name: "sqlite_scan".into(),
execute: scan_fn,
},
);
tracing::info!("SQLite extension loaded: 2 functions registered (rusqlite native)");
}
#[cfg(not(feature = "bundled"))]
{
context.register_scalar_function(
"sqlite_query",
ScalarFunction::CustomScalar {
name: "sqlite_query".into(),
execute: Arc::new(|_| Err("SQLite not available (feature 'bundled' disabled)".into())),
},
);
context.register_table_function(
"sqlite_scan",
TableFunction::CustomTable {
name: "sqlite_scan".into(),
execute: Arc::new(|_, _| Err("SQLite not available (feature 'bundled' disabled)".into())),
},
);
tracing::info!("SQLite extension loaded: 2 functions registered (placeholder)");
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sqlite_extension_name() {
let ext = SqliteExtension::new();
assert_eq!(ext.name(), "SQLITE");
}
}