diff --git a/src/db/general.rs b/src/db/general.rs index 940de21..ae01c57 100644 --- a/src/db/general.rs +++ b/src/db/general.rs @@ -30,6 +30,18 @@ impl Account { Ok(account) } + pub async fn get_or_create(db: &Surreal, address: &str) -> Result { + let id = (ACCOUNT, address); + + match db.select(id).await? { + Some(account) => Ok(account), + None => { + let account: Option = db.create(id).await?; + account.ok_or(Error::FailedToCreateDBEntry) + } + } + } + pub async fn airdrop(db: &Surreal, account: &str, tokens: u64) -> Result<(), Error> { let tokens = tokens.saturating_mul(1_000_000_000); let _ = db diff --git a/src/db/mod.rs b/src/db/mod.rs index 5e2e5f4..61c510a 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -22,6 +22,8 @@ pub enum Error { StdIo(#[from] std::io::Error), #[error(transparent)] TimeOut(#[from] tokio::time::error::Elapsed), + #[error("Failed to create account")] + FailedToCreateDBEntry, } pub mod prelude { diff --git a/src/db/vm.rs b/src/db/vm.rs index 8d7ca71..6f4c8fb 100644 --- a/src/db/vm.rs +++ b/src/db/vm.rs @@ -3,6 +3,7 @@ use std::time::Duration; use super::Error; use crate::constants::{ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, VM_NODE}; +use crate::db::general; use crate::db::general::Report; use crate::old_brain; use serde::{Deserialize, Serialize}; @@ -50,6 +51,7 @@ impl VmNodeResources { impl VmNode { pub async fn register(self, db: &Surreal) -> Result<(), Error> { + general::Account::get_or_create(db, &self.operator.key().to_string()).await?; let _: Option = db.upsert(self.id.clone()).content(self).await?; Ok(()) }