From 18c647ff3e64a6e646f8087aa52dc1c593773d20 Mon Sep 17 00:00:00 2001
From: Noor <noormohammedb@protonmail.com>
Date: Wed, 7 May 2025 20:33:51 +0530
Subject: [PATCH] fix: register vm node creates operator account in db

---
 src/db/general.rs | 12 ++++++++++++
 src/db/mod.rs     |  2 ++
 src/db/vm.rs      |  3 ++-
 3 files changed, 16 insertions(+), 1 deletion(-)

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<Client>, address: &str) -> Result<Self, Error> {
+        let id = (ACCOUNT, address);
+
+        match db.select(id).await? {
+            Some(account) => Ok(account),
+            None => {
+                let account: Option<Self> = db.create(id).await?;
+                account.ok_or(Error::FailedToCreateDBEntry)
+            }
+        }
+    }
+
     pub async fn airdrop(db: &Surreal<Client>, 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 24d08ee..f6968c4 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 6b87bfc..2da35e4 100644
--- a/src/db/vm.rs
+++ b/src/db/vm.rs
@@ -5,7 +5,7 @@ use super::Error;
 use crate::constants::{
     ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, UPDATE_VM_REQ, VM_NODE, VM_UPDATE_EVENT,
 };
-use crate::db::{MeasurementArgs, Report, VmNodeFilters};
+use crate::db::{Account, MeasurementArgs, Report, VmNodeFilters};
 use crate::old_brain;
 use serde::{Deserialize, Serialize};
 use surrealdb::engine::remote::ws::Client;
@@ -52,6 +52,7 @@ impl VmNodeResources {
 
 impl VmNode {
     pub async fn register(self, db: &Surreal<Client>) -> Result<(), Error> {
+        Account::get_or_create(db, &self.operator.key().to_string()).await?;
         let _: Option<VmNode> = db.upsert(self.id.clone()).content(self).await?;
         Ok(())
     }