brain/src/db/mod.rs
Noor 4dfaa3f465
fix: type matching on live_vmnode_msgs
generic type matching to table name
update type paths with
some logging in tests
2025-05-06 17:25:42 +05:30

119 lines
3.8 KiB
Rust

pub mod app;
pub mod general;
pub mod vm;
use crate::constants::{DELETED_VM, NEW_VM_REQ, UPDATE_VM_REQ};
use crate::old_brain;
use prelude::*;
use serde::{Deserialize, Serialize};
use surrealdb::engine::remote::ws::{Client, Ws};
use surrealdb::opt::auth::Root;
use surrealdb::{Notification, Surreal};
use tokio::sync::mpsc::Sender;
use tokio_stream::StreamExt;
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("Internal DB error: {0}")]
DataBase(#[from] surrealdb::Error),
#[error("Daemon channel got closed: {0}")]
VmDaemonConnection(#[from] tokio::sync::mpsc::error::SendError<VmDaemonMsg>),
#[error(transparent)]
StdIo(#[from] std::io::Error),
#[error(transparent)]
TimeOut(#[from] tokio::time::error::Elapsed),
}
pub mod prelude {
pub use super::app::*;
pub use super::general::*;
pub use super::vm::*;
pub use super::*;
}
pub async fn db_connection(
db_address: &str,
username: &str,
password: &str,
ns: &str,
db: &str,
) -> Result<Surreal<Client>, Error> {
let db_connection: Surreal<Client> = Surreal::init();
db_connection.connect::<Ws>(db_address).await?;
// Sign in to the server
db_connection.signin(Root { username, password }).await?;
db_connection.use_ns(ns).use_db(db).await?;
Ok(db_connection)
}
pub async fn migration0(
db: &Surreal<Client>,
old_data: &old_brain::BrainData,
) -> Result<(), Error> {
let accounts: Vec<Account> = old_data.into();
let vm_nodes: Vec<VmNode> = old_data.into();
let app_nodes: Vec<AppNode> = old_data.into();
let vm_contracts: Vec<ActiveVm> = old_data.into();
let schema = std::fs::read_to_string(crate::constants::DB_SCHEMA_FILE)?;
db.query(schema).await?;
println!("Inserting accounts...");
let _: Vec<Account> = db.insert(()).content(accounts).await?;
println!("Inserting vm nodes...");
let _: Vec<VmNode> = db.insert(()).content(vm_nodes).await?;
println!("Inserting app nodes...");
let _: Vec<AppNode> = db.insert(()).content(app_nodes).await?;
println!("Inserting vm contracts...");
let _: Vec<ActiveVm> = db.insert("vm_contract").relation(vm_contracts).await?;
Ok(())
}
pub async fn upsert_record<SomeRecord: Serialize + 'static>(
db: &Surreal<Client>,
table: &str,
id: &str,
my_record: SomeRecord,
) -> Result<(), Error> {
#[derive(Deserialize)]
struct Wrapper {}
let _: Option<Wrapper> = db.create((table, id)).content(my_record).await?;
Ok(())
}
pub async fn live_vmnode_msgs<
T: Into<vm::VmDaemonMsg> + std::marker::Unpin + for<'de> Deserialize<'de>,
>(
db: &Surreal<Client>,
node: &str,
tx: Sender<vm::VmDaemonMsg>,
) -> Result<(), Error> {
let table_name = match std::any::type_name::<T>() {
t if t == std::any::type_name::<crate::db::vm::NewVmReq>() => NEW_VM_REQ.to_string(),
t if t == std::any::type_name::<crate::db::vm::UpdateVmReq>() => UPDATE_VM_REQ.to_string(),
t if t == std::any::type_name::<crate::db::vm::DeletedVm>() => DELETED_VM.to_string(),
wat => {
log::error!("listen_for_node: T has type {wat}");
String::from("wat")
}
};
let mut resp =
db.query(format!("live select * from {table_name} where out = vm_node:{node};")).await?;
let mut live_stream = resp.stream::<Notification<T>>(0)?;
while let Some(result) = live_stream.next().await {
match result {
Ok(notification) => {
if notification.action == surrealdb::Action::Create {
tx.send(notification.data.into()).await?
}
}
Err(e) => {
log::error!("listen_for_{table_name} DB stream failed for {node}: {e}");
return Err(Error::from(e));
}
}
}
Ok(())
}