use std::sync::Arc;
use aetherdb::{ClusterReconfigCmd, ControlPlaneCmd, DmaRequest, Key};
use kdl::KdlNode;
use crate::aether_host::host_config_from_parts;
use crate::bus::{response, BusModule};
use crate::kwt_replay::runtime::{
format_cluster_response, format_dma_response, safety_config_from_kdl, AetherRuntime,
RuntimeError,
};
const DEFAULT_MIGRATE_SESSION: &str = "default";
pub struct AetherRuntimeModule {
runtime: Arc<AetherRuntime>,
}
impl AetherRuntimeModule {
#[must_use]
pub fn new(runtime: Arc<AetherRuntime>) -> Self {
Self { 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 require_admin() -> Result<(), String> {
if crate::surfaces::env_enabled("NAUTALID_AETHERDB_BUS_MUTATE")
|| crate::surfaces::env_enabled("NAUTALID_AETHERDB_BUS_ADMIN")
{
Ok(())
} else {
Err("runtime admin ops require NAUTALID_AETHERDB_BUS_ADMIN=1 or NAUTALID_AETHERDB_BUS_MUTATE=1".into())
}
}
fn opt_string(node: &KdlNode, name: &str) -> Option<String> {
node.get(name)
.and_then(|v| v.as_string())
.map(str::to_string)
}
fn req_string(node: &KdlNode, name: &str) -> Result<String, String> {
Self::opt_string(node, name)
.filter(|s| !s.is_empty())
.ok_or_else(|| format!("missing `{name}`"))
}
fn opt_u32(node: &KdlNode, name: &str) -> Option<u32> {
node.get(name)
.and_then(|v| v.as_integer())
.and_then(|n| u32::try_from(n).ok())
}
fn opt_u64(node: &KdlNode, name: &str) -> Option<u64> {
node.get(name)
.and_then(|v| v.as_integer())
.and_then(|n| u64::try_from(n).ok())
}
fn subcommand_node<'a>(cmd: &'a KdlNode) -> Option<&'a KdlNode> {
cmd.children()?.nodes().first()
}
fn map_err(err: RuntimeError) -> String {
response::err(&err.to_string())
}
fn migrate_session(cmd: &KdlNode) -> String {
Self::opt_string(cmd, "session").unwrap_or_else(|| DEFAULT_MIGRATE_SESSION.into())
}
}
impl BusModule for AetherRuntimeModule {
fn name(&self) -> &'static str {
"aether-runtime"
}
fn handle(&self, module_node: &KdlNode, command_nodes: &[&KdlNode]) -> String {
let Some(cmd) = Self::command_node(module_node, command_nodes) else {
return response::err(
"aether-runtime commands: status, apply, host, jti, spawn, close, cluster, control, catalog, lifecycle, databases, dma, checkpoint, emergency-persist, write-permit, short-circuit",
);
};
match cmd.name().value() {
"status" => {
let topo = self.runtime.topology();
let host = self.runtime.host_status();
format!(
"status mode=\"{}\" shards={} replication={} node_id={} query=\"{}\" host_running={} host_endpoint=\"{}\" primary_id=\"{}\"",
topo.mode_label(),
topo.shards,
topo.replication,
topo.node_id,
topo.query.replace('"', "\\\""),
host.running,
host.endpoint.unwrap_or_default().replace('"', "\\\""),
host.primary_id.unwrap_or_default().replace('"', "\\\""),
)
}
"apply" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
match self.runtime.apply_topology_from_kdl(cmd) {
Ok(()) => response::ok("topology applied"),
Err(err) => Self::map_err(err),
}
}
"host" => self.handle_host(cmd),
"jti" => self.handle_jti(cmd),
"spawn" => self.handle_spawn(cmd),
"close" => self.handle_close(cmd),
"cluster" => self.handle_cluster(cmd),
"control" => self.handle_control(cmd),
"catalog" => match self.runtime.database_catalog_kdl() {
Ok(kdl) => format!("ok catalog\n{kdl}"),
Err(RuntimeError::ClusterRequired) => {
response::err("catalog requires cluster data plane")
}
Err(RuntimeError::HostStopped) => response::err("data plane not running"),
Err(err) => Self::map_err(err),
},
"lifecycle" => match self.runtime.lifecycle_kdl() {
Ok(kdl) => format!("ok lifecycle\n{kdl}"),
Err(RuntimeError::ClusterRequired) => {
response::err("lifecycle requires cluster data plane")
}
Err(RuntimeError::HostStopped) => response::err("data plane not running"),
Err(err) => Self::map_err(err),
},
"databases" => match self.runtime.database_addresses() {
Ok(addresses) => format!("ok databases {}", addresses.join(" ")),
Err(RuntimeError::ClusterRequired) => {
response::err("databases requires cluster data plane")
}
Err(RuntimeError::HostStopped) => response::err("data plane not running"),
Err(err) => Self::map_err(err),
},
"dma" => self.handle_dma(cmd),
"checkpoint" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let address = match Self::req_string(cmd, "address") {
Ok(a) => a,
Err(e) => return response::err(&e),
};
match self.runtime.checkpoint_at(&address) {
Ok(seq) => format!("ok checkpoint seq={seq}"),
Err(err) => Self::map_err(err),
}
}
"emergency-persist" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let address = match Self::req_string(cmd, "address") {
Ok(a) => a,
Err(e) => return response::err(&e),
};
match self.runtime.emergency_persist_at(&address) {
Ok((path, seq)) => format!(
"ok emergency_persist path=\"{}\" seq={seq}",
path.display().to_string().replace('"', "\\\"")
),
Err(err) => Self::map_err(err),
}
}
"write-permit" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
match self.runtime.issue_write_permit() {
Ok(permit) => format!(
"ok write_permit unrestricted={}",
permit.is_unrestricted()
),
Err(err) => Self::map_err(err),
}
}
"short-circuit" => match self.runtime.short_circuit_dump() {
Ok(report) => format!("ok short_circuit databases={}", report.databases.len()),
Err(RuntimeError::ClusterRequired) => {
response::err("short-circuit requires cluster data plane")
}
Err(RuntimeError::HostStopped) => response::err("data plane not running"),
Err(err) => Self::map_err(err),
},
other => response::err(&format!("unknown aether-runtime command: {other}")),
}
}
}
impl AetherRuntimeModule {
fn handle_spawn(&self, cmd: &KdlNode) -> String {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let id = match Self::req_string(cmd, "id") {
Ok(id) => id,
Err(e) => return response::err(&e),
};
let preset = Self::opt_string(cmd, "preset").unwrap_or_else(|| "nautalid-data".into());
let source = match Self::opt_string(cmd, "kdl") {
Some(kdl) => kdl,
None => return response::err("spawn requires kdl=\"…\""),
};
match self.runtime.spawn_database(&id, &source, &preset) {
Ok(db) => format!("ok spawned id=\"{id}\" address=\"{}\"", db.address()),
Err(err) => Self::map_err(err),
}
}
fn handle_close(&self, cmd: &KdlNode) -> String {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let id = match Self::req_string(cmd, "id") {
Ok(id) => id,
Err(e) => return response::err(&e),
};
match self.runtime.close_database(&id) {
Ok(()) => response::ok(&format!("closed {id}")),
Err(err) => Self::map_err(err),
}
}
fn handle_host(&self, cmd: &KdlNode) -> String {
let sub = Self::subcommand_node(cmd).map(|n| n.name().value()).unwrap_or("");
match sub {
"start" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let bind = Self::opt_string(cmd, "bind");
let replicate_key_hex = Self::opt_string(cmd, "replicate_key_hex")
.or_else(|| Self::opt_string(cmd, "replicate_hex"));
let control_key_hex = Self::opt_string(cmd, "control_key_hex");
let cfg = match host_config_from_parts(
bind,
Self::opt_string(cmd, "data_dir"),
replicate_key_hex,
control_key_hex,
) {
Ok(cfg) => cfg,
Err(e) => return response::err(&e),
};
let kdl = Self::opt_string(cmd, "kdl");
match self.runtime.start_host(cfg, kdl) {
Ok(()) => response::ok("host started"),
Err(err) => Self::map_err(err),
}
}
"stop" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
match self.runtime.stop_host() {
Ok(()) => response::ok("host stopped"),
Err(err) => Self::map_err(err),
}
}
"status" => {
let host = self.runtime.host_status();
format!(
"status running={} mode=\"{}\" endpoint=\"{}\"",
host.running,
host.mode.replace('"', "\\\""),
host.endpoint.unwrap_or_default().replace('"', "\\\"")
)
}
_ => response::err("host commands: start, stop, status"),
}
}
fn handle_jti(&self, cmd: &KdlNode) -> String {
let sub = Self::subcommand_node(cmd).map(|n| n.name().value()).unwrap_or("");
match sub {
"reconnect" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let url = match Self::req_string(cmd, "url") {
Ok(url) => url,
Err(e) => return response::err(&e),
};
match self.runtime.jti_reconnect(&url) {
Ok(()) => response::ok("jti reconnected"),
Err(err) => Self::map_err(err),
}
}
"attach" => {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
match self.runtime.jti_attach_primary() {
Ok(()) => response::ok("jti attached to primary"),
Err(err) => Self::map_err(err),
}
}
_ => response::err("jti commands: reconnect, attach"),
}
}
fn handle_dma(&self, cmd: &KdlNode) -> String {
let address = match Self::req_string(cmd, "address") {
Ok(a) => a,
Err(e) => return response::err(&e),
};
let active = Self::subcommand_node(cmd).unwrap_or(cmd);
let sub = active.name().value();
let request = match sub {
"get" => {
let key = match active.get("key").and_then(|v| v.as_string()) {
Some(k) => Key::from_bytes(k.as_bytes().to_vec()),
None => return response::err("dma get requires key"),
};
DmaRequest::Get(key)
}
"range" => {
let from = match active.get("from").and_then(|v| v.as_string()) {
Some(k) => Key::from_bytes(k.as_bytes().to_vec()),
None => return response::err("dma range requires from"),
};
let to = match active.get("to").and_then(|v| v.as_string()) {
Some(k) => Key::from_bytes(k.as_bytes().to_vec()),
None => return response::err("dma range requires to"),
};
DmaRequest::Range { from, to }
}
"sql" => {
let text = match Self::req_string(active, "text") {
Ok(t) => t,
Err(e) => return response::err(&e),
};
DmaRequest::SqlSelect(text)
}
"hybrid" => {
let text = match Self::req_string(active, "text") {
Ok(t) => t,
Err(e) => return response::err(&e),
};
DmaRequest::HybridSelect(text)
}
_ => return response::err("dma commands: get, range, sql, hybrid"),
};
format_dma_response(&self.runtime.dma_at(&address, request))
}
fn handle_cluster(&self, cmd: &KdlNode) -> String {
let sub = Self::subcommand_node(cmd)
.map(|n| n.name().value())
.unwrap_or("status");
if sub != "status" && sub != "short-circuit" {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
}
let reconfig = match sub {
"status" => ClusterReconfigCmd::GetStatus,
"set-read-only" => ClusterReconfigCmd::SetReadOnly {
enabled: cmd
.get("enabled")
.and_then(|v| v.as_bool())
.or_else(|| {
cmd.get("enabled")
.and_then(|v| v.as_string())
.map(|s| matches!(s, "1" | "true" | "yes"))
})
.unwrap_or(true),
},
"set-shards" => ClusterReconfigCmd::SetShards {
count: Self::opt_u32(cmd, "count").unwrap_or(1).max(1),
},
"apply-schema" => {
let kdl = match Self::req_string(cmd, "kdl") {
Ok(kdl) => kdl,
Err(e) => return response::err(&e),
};
ClusterReconfigCmd::ApplySchema { kdl }
}
"set-safety" => match safety_config_from_kdl(cmd) {
Ok(cfg) => ClusterReconfigCmd::SetSafety(cfg),
Err(err) => return Self::map_err(err),
},
_ => {
return response::err(
"cluster commands: status, set-read-only, set-shards, apply-schema, set-safety",
);
}
};
format_cluster_response(&self.runtime.cluster_reconfig(reconfig))
}
fn handle_control(&self, cmd: &KdlNode) -> String {
let sub_node = Self::subcommand_node(cmd);
let sub = sub_node.map(|n| n.name().value()).unwrap_or("");
if sub == "forward" {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let forward_node = sub_node.expect("forward subcommand");
let target = Self::opt_u64(forward_node, "target_node_id").unwrap_or(0);
if target == 0 {
return response::err("forward requires target_node_id");
}
let Some(inner) = forward_node.children().and_then(|c| c.nodes().first()) else {
return response::err("forward requires nested control command");
};
let plane = match Self::control_plane_from_node(inner) {
Ok(plane) => plane,
Err(e) => return response::err(&e),
};
return format_cluster_response(&self.runtime.forward_control_cmd(target, plane));
}
let active = sub_node.unwrap_or(cmd);
let sub = active.name().value();
if sub == "migrate-import" || sub == "migrate-import-bundle" {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
let session = Self::migrate_session(active);
let database_id = Self::opt_string(active, "database_id");
return format_cluster_response(
&self
.runtime
.migrate_import_cached(&session, database_id),
);
}
if sub != "placement" && sub != "balance" {
if let Err(e) = Self::require_admin() {
return response::err(&e);
}
}
let session = Self::migrate_session(active);
let database_id = Self::opt_string(active, "database_id").unwrap_or_default();
let plane = match Self::control_plane_from_node(active) {
Ok(plane) => plane,
Err(e) => return response::err(&e),
};
let resp = self.runtime.control_plane(plane);
if matches!(sub, "migrate-export" | "migrate-export-bundle") && !database_id.is_empty() {
self.runtime
.store_migrate_export(&session, &database_id, &resp);
let mut body = format_cluster_response(&resp);
body.push_str(&format!(" session=\"{session}\""));
return body;
}
format_cluster_response(&resp)
}
fn control_plane_from_node(node: &KdlNode) -> Result<ControlPlaneCmd, String> {
let sub = node.name().value();
Ok(match sub {
"register-node" => ControlPlaneCmd::RegisterNode {
node_id: Self::opt_u64(node, "node_id").unwrap_or(1),
endpoint: Self::opt_string(node, "endpoint").unwrap_or_default(),
},
"unregister-node" => ControlPlaneCmd::UnregisterNode {
node_id: Self::opt_u64(node, "node_id").unwrap_or(0),
},
"assign-shard" => ControlPlaneCmd::AssignShard {
shard_id: Self::opt_u32(node, "shard_id").unwrap_or(0),
node_id: Self::opt_u64(node, "node_id").unwrap_or(0),
},
"placement" => ControlPlaneCmd::GetPlacement,
"balance" => ControlPlaneCmd::BalanceReport,
"migrate-export" => ControlPlaneCmd::MigrateExport {
database_id: Self::opt_string(node, "database_id").unwrap_or_default(),
logical_shard: Self::opt_u32(node, "logical_shard").unwrap_or(0),
},
"migrate-export-bundle" => ControlPlaneCmd::MigrateExportBundle {
database_id: Self::opt_string(node, "database_id").unwrap_or_default(),
logical_shard: Self::opt_u32(node, "logical_shard").unwrap_or(0),
},
"migrate-import" => ControlPlaneCmd::MigrateImport {
database_id: Self::opt_string(node, "database_id").unwrap_or_default(),
rows: Vec::new(),
},
"migrate-import-bundle" => {
return Err("use control migrate-import-bundle with session= from a prior export".into());
}
"balance-execute" => ControlPlaneCmd::BalanceExecute {
local_node_id: Self::opt_u64(node, "local_node_id").unwrap_or(1),
},
other => {
return Err(format!(
"unknown control subcommand `{other}` (try forward for remote ops)"
));
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use kdl::KdlDocument;
#[test]
fn status_without_host() {
let runtime = Arc::new(AetherRuntime::bootstrap_from_env().expect("bootstrap"));
let doc: KdlDocument = "aether-runtime {\n status\n}".parse().unwrap();
let node = doc.nodes().first().unwrap();
let body = AetherRuntimeModule::new(runtime).handle(node, &[]);
assert!(body.contains("host_running=false"));
}
}