brain/src/db/general.rs

use crate::constants::ACCOUNT;
use crate::db::prelude::*;

use super::Error;
use crate::old_brain;
use serde::{Deserialize, Serialize};
use surrealdb::engine::remote::ws::Client;
use surrealdb::sql::Datetime;
use surrealdb::{RecordId, Surreal};

#[derive(Debug, Serialize, Deserialize)]
pub struct Account {
    pub id: RecordId,
    pub balance: u64,
    pub tmp_locked: u64,
    pub escrow: u64,
    pub email: String,
}

impl Account {
    pub async fn get(db: &Surreal<Client>, address: &str) -> Result<Self, Error> {
        let id = (ACCOUNT, address);
        let account: Option<Self> = db.select(id).await?;
        let account = match account {
            Some(account) => account,
            None => {
                Self { id: id.into(), balance: 0, tmp_locked: 0, escrow: 0, email: String::new() }
            }
        };
        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
            .query(format!("upsert account:{account} SET balance = (balance || 0) + {tokens};"))
            .await?;
        Ok(())
    }
}

impl Account {
    pub async fn is_banned_by_node(
        db: &Surreal<Client>,
        user: &str,
        node: &str,
    ) -> Result<bool, Error> {
        let ban: Option<Self> = db
            .query(format!(
                "(select operator->ban[0] as ban
                    from vm_node:{node}
                    where operator->ban->account contains account:{user}
                ).ban;"
            ))
            .await?
            .take(0)?;
        Ok(ban.is_some())
    }
}

impl From<&old_brain::BrainData> for Vec<Account> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut accounts = Vec::new();
        for old_account in old_data.accounts.iter() {
            let mut a = Account {
                id: RecordId::from(("account", old_account.key())),
                balance: old_account.value().balance,
                tmp_locked: old_account.value().tmp_locked,
                escrow: 0,
                email: String::new(),
            };
            if let Some(operator) = old_data.operators.get(old_account.key()) {
                a.escrow = operator.escrow;
                a.email = operator.email.clone();
            }
            accounts.push(a);
        }
        accounts
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Ban {
    id: RecordId,
    #[serde(rename = "in")]
    from_account: RecordId,
    #[serde(rename = "out")]
    to_account: RecordId,
    created_at: Datetime,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Kick {
    id: RecordId,
    #[serde(rename = "in")]
    from_account: RecordId,
    #[serde(rename = "out")]
    to_account: RecordId,
    created_at: Datetime,
    reason: String,
    contract: RecordId,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Report {
    #[serde(rename = "in")]
    from_account: RecordId,
    #[serde(rename = "out")]
    to_node: RecordId,
    created_at: Datetime,
    pub reason: String,
    pub contract_id: String,
}

impl Report {
    // TODO: test this functionality and remove this comment
    pub async fn create(
        db: &Surreal<Client>,
        from_account: RecordId,
        to_node: RecordId,
        reason: String,
        contract_id: String,
    ) -> Result<(), Error> {
        let _: Vec<Self> = db
            .insert("report")
            .relation(Report {
                from_account,
                to_node,
                created_at: Datetime::default(),
                reason,
                contract_id,
            })
            .await?;
        Ok(())
    }
}

/// This is the operator obtained from the DB,
/// however the relation is defined using OperatorRelation
#[derive(Debug, Serialize, Deserialize)]
pub struct Operator {
    pub account: RecordId,
    pub app_nodes: u64,
    pub vm_nodes: u64,
    pub email: String,
    pub escrow: u64,
    pub reports: u64,
}

impl Operator {
    pub async fn list(db: &Surreal<Client>) -> Result<Vec<Self>, Error> {
        let mut result = db
            .query(
                "array::distinct(array::flatten( [
                    (select operator from vm_node group by operator).operator,
                    (select operator from app_node group by operator).operator
                ]));"
                    .to_string(),
            )
            .await?;
        let operator_accounts: Vec<RecordId> = result.take(0)?;
        let mut operators: Vec<Self> = Vec::new();
        for account in operator_accounts.iter() {
            if let Some(operator) = Self::inspect(db, &account.key().to_string()).await? {
                operators.push(operator);
            }
        }
        Ok(operators)
    }

    pub async fn inspect(db: &Surreal<Client>, account: &str) -> Result<Option<Self>, Error> {
        let mut result = db
            .query(format!(
                "$vm_nodes = (select id from vm_node where operator = account:{account}).id;
                $app_nodes = (select id from app_node where operator = account:{account}).id;
                select *,
                    id as account,
                    email,
                    escrow,
                    $vm_nodes.len() as vm_nodes,
                    $app_nodes.len() as app_nodes,
                    (select id from report where $vm_nodes contains out).len() +
                    (select id from report where $app_nodes contains out).len()
                    as reports
                from account where id = account:{account};"
            ))
            .await?;
        let operator: Option<Self> = result.take(2)?;
        Ok(operator)
    }

    pub async fn inspect_nodes(
        db: &Surreal<Client>,
        account: &str,
    ) -> Result<(Option<Self>, Vec<VmNodeWithReports>, Vec<AppNodeWithReports>), Error> {
        let operator = Self::inspect(db, account).await?;
        let mut result = db
            .query(format!(
                "select *, operator, <-report.* as reports from vm_node
                 where operator = account:{account};"
            ))
            .query(format!(
                "select *, operator, <-report.* as reports from app_node
                 where operator = account:{account};"
            ))
            .await?;
        let vm_nodes: Vec<VmNodeWithReports> = result.take(0)?;
        let app_nodes: Vec<AppNodeWithReports> = result.take(1)?;

        Ok((operator, vm_nodes, app_nodes))
    }
}