use std::sync::Arc;
use kdl::KdlNode;
use crate::bus::{response, BusModule};
use crate::kwt_replay::bus_api::{self, BusDbError};
use crate::kwt_replay::JtiStore;
pub struct AetherDbModule {
store: Arc<JtiStore>,
runtime: Option<Arc<crate::kwt_replay::AetherRuntime>>,
}
impl AetherDbModule {
#[must_use]
pub fn new(store: Arc<JtiStore>) -> Self {
Self {
store,
runtime: None,
}
}
#[must_use]
pub fn with_runtime(store: Arc<JtiStore>, runtime: Arc<crate::kwt_replay::AetherRuntime>) -> Self {
Self {
store,
runtime: Some(runtime),
}
}
fn command_node<'a>(
module_node: &'a KdlNode,
command_nodes: &[&'a KdlNode],
) -> Option<&'a KdlNode> {
if let Some(cmd) = command_nodes.first() {
return Some(cmd);
}
module_node.children()?.nodes().first()
}
fn required_bytes(node: &KdlNode, name: &str) -> Result<Vec<u8>, BusDbError> {
let raw = node
.get(name)
.and_then(|value| value.as_string())
.ok_or_else(|| BusDbError::MissingProperty(name.to_string()))?;
Ok(raw.as_bytes().to_vec())
}
fn optional_text(node: &KdlNode, name: &str) -> Result<String, BusDbError> {
node.get(name)
.and_then(|value| value.as_string())
.map(str::to_string)
.ok_or_else(|| BusDbError::MissingProperty(name.to_string()))
}
fn map_err(err: BusDbError) -> String {
response::err(&err.to_string())
}
}
impl BusModule for AetherDbModule {
fn name(&self) -> &'static str {
"aetherdb"
}
fn handle(&self, module_node: &KdlNode, command_nodes: &[&KdlNode]) -> String {
let Some(cmd) = Self::command_node(module_node, command_nodes) else {
return response::err(
"aetherdb commands: status, capabilities, get, put, delete, query, query-at, sql, cypher",
);
};
match cmd.name().value() {
"status" => {
let topology = self
.runtime
.as_ref()
.map(|r| r.topology())
.unwrap_or_else(crate::kwt_replay::Topology::from_env);
match self.store.bus_status_with_topology(&topology) {
Ok(status) => format!(
"status enabled={} shared={} mode=\"{}\" shards={} replication={} node_id={} jti_prefix=\"{}\" blades=\"{}\"",
status.enabled,
status.shared,
status.mode,
status.shards,
status.replication,
status.node_id,
status.jti_prefix.replace('"', "\\\""),
status.blades.join(",")
),
Err(BusDbError::Disabled) => {
"status enabled=false shared=false mode=\"local\" shards=0 replication=0 node_id=0 jti_prefix=\"\" blades=\"\"".into()
}
Err(err) => Self::map_err(err),
}
}
"capabilities" => {
let topology = self
.runtime
.as_ref()
.map(|r| r.topology())
.unwrap_or_else(crate::kwt_replay::Topology::from_env);
match self.store.bus_status_with_topology(&topology) {
Ok(status) => format!("capabilities {}", status.blades.join(" ")),
Err(BusDbError::Disabled) => "capabilities".into(),
Err(err) => Self::map_err(err),
}
}
"get" => match Self::required_bytes(cmd, "key").and_then(|key| self.store.bus_get(&key)) {
Ok(Some(value)) => {
format!("ok hit {}", bus_api::format_bytes_attr("value", &value))
}
Ok(None) => "ok miss".into(),
Err(err) => Self::map_err(err),
},
"put" => {
let key = Self::required_bytes(cmd, "key");
let value = Self::required_bytes(cmd, "value");
match key.and_then(|key| {
value.and_then(|value| self.store.bus_put(&key, &value).map(|()| key))
}) {
Ok(key) => format!("ok put {}", bus_api::format_bytes_attr("key", &key)),
Err(err) => Self::map_err(err),
}
}
"delete" => {
match Self::required_bytes(cmd, "key").and_then(|key| self.store.bus_delete(&key)) {
Ok(()) => response::ok("deleted"),
Err(err) => Self::map_err(err),
}
}
"query" => {
match Self::optional_text(cmd, "text").and_then(|text| self.store.bus_query(&text))
{
Ok(body) => body,
Err(err) => Self::map_err(err),
}
}
"sql" => {
match Self::optional_text(cmd, "text").and_then(|text| self.store.bus_sql(&text)) {
Ok(body) => body,
Err(err) => Self::map_err(err),
}
}
"cypher" => {
match Self::optional_text(cmd, "text").and_then(|text| self.store.bus_cypher(&text))
{
Ok(body) => body,
Err(err) => Self::map_err(err),
}
}
"query-at" => {
let address = match Self::optional_text(cmd, "address") {
Ok(a) => a,
Err(err) => return Self::map_err(err),
};
let text = match Self::optional_text(cmd, "text") {
Ok(t) => t,
Err(err) => return Self::map_err(err),
};
let Some(runtime) = &self.runtime else {
return response::err("query-at requires unified LID runtime");
};
match runtime.query_at(&address, &text) {
Ok(result) => bus_api::format_query_result_public(&result),
Err(err) => response::err(&err.to_string()),
}
}
other => response::err(&format!("unknown aetherdb command: {other}")),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use kdl::KdlDocument;
#[test]
fn disabled_status() {
let store = Arc::new(JtiStore::disabled());
let doc: KdlDocument = "aetherdb {\n status\n}".parse().unwrap();
let node = doc.nodes().first().unwrap();
let body = AetherDbModule::new(store).handle(node, &[]);
assert!(body.contains("enabled=false"));
}
#[test]
fn get_miss_without_key_errors() {
let store = Arc::new(JtiStore::disabled());
let doc: KdlDocument = "aetherdb {\n get\n}".parse().unwrap();
let node = doc.nodes().first().unwrap();
let body = AetherDbModule::new(store).handle(node, &[]);
assert!(body.starts_with("err "));
}
}