brain-to-surreal/src/db.rs

use crate::old_brain;
use serde::{Deserialize, Serialize};
use std::sync::LazyLock;
use surrealdb::{
    engine::remote::ws::{Client, Ws},
    opt::auth::Root,
    sql::Datetime,
    RecordId, Surreal,
};

static DB: LazyLock<Surreal<Client>> = LazyLock::new(Surreal::init);
const ACCOUNT: &str = "account";
const OPERATOR: &str = "operator";
const VM_CONTRACT: &str = "vm_contract";
const VM_NODE: &str = "vm_node";

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error(transparent)]
    DataBase(#[from] surrealdb::Error),
}

pub async fn init() -> surrealdb::Result<()> {
    DB.connect::<Ws>("localhost:8000").await?;
    // Sign in to the server
    DB.signin(Root { username: "root", password: "root" }).await?;
    DB.use_ns("brain").use_db("migration").await?;
    Ok(())
}

pub async fn migration0(old_data: &old_brain::BrainData) -> surrealdb::Result<()> {
    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<VmContract> = old_data.into();
    let operators: Vec<OperatorRelation> = old_data.into();
    let app_contracts: Vec<AppContract> = old_data.into();

    init().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<VmContract> = DB.insert("vm_contract").relation(vm_contracts).await?;
    println!("Inserting app contracts...");
    let _: Vec<AppContract> = DB.insert("app_contract").relation(app_contracts).await?;
    println!("Inserting operators...");
    let _: Vec<OperatorRelation> = DB.insert(OPERATOR).relation(operators).await?;

    Ok(())
}

#[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(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 airdrop(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(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct VmNode {
    pub id: RecordId,
    pub country: String,
    pub region: String,
    pub city: String,
    pub ip: String,
    pub avail_mem_mb: u32,
    pub avail_vcpus: u32,
    pub avail_storage_gbs: u32,
    pub avail_ipv4: u32,
    pub avail_ipv6: u32,
    pub avail_ports: u32,
    pub max_ports_per_vm: u32,
    pub price: u64,
    pub offline_minutes: u64,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct VmContract {
    pub id: RecordId,
    #[serde(rename = "in")]
    pub admin: RecordId,
    #[serde(rename = "out")]
    pub vm_node: RecordId,
    pub state: String,
    pub hostname: String,
    pub mapped_ports: Vec<(u32, u32)>,
    pub public_ipv4: String,
    pub public_ipv6: String,
    pub disk_size_gb: u32,
    pub vcpus: u32,
    pub memory_mb: u32,
    pub dtrfs_sha: String,
    pub kernel_sha: String,
    pub created_at: Datetime,
    pub updated_at: Datetime,
    pub price_per_unit: u64,
    pub locked_nano: u64,
    pub collected_at: Datetime,
}

impl VmContract {
    /// total hardware units of this VM
    fn total_units(&self) -> u64 {
        // TODO: Optimize this based on price of hardware.
        // I tried, but this can be done better.
        // Storage cost should also be based on tier
        (self.vcpus as u64 * 10)
            + ((self.memory_mb + 256) as u64 / 200)
            + (self.disk_size_gb as u64 / 10)
            + (!self.public_ipv4.is_empty() as u64 * 10)
    }

    /// Returns price per minute in nanoLP
    pub fn price_per_minute(&self) -> u64 {
        self.total_units() * self.price_per_unit
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct VmContractWithNode {
    pub id: RecordId,
    #[serde(rename = "in")]
    pub admin: RecordId,
    #[serde(rename = "out")]
    pub vm_node: VmNode,
    pub state: String,
    pub hostname: String,
    pub mapped_ports: Vec<(u32, u32)>,
    pub public_ipv4: String,
    pub public_ipv6: String,
    pub disk_size_gb: u32,
    pub vcpus: u32,
    pub memory_mb: u32,
    pub dtrfs_sha: String,
    pub kernel_sha: String,
    pub created_at: Datetime,
    pub updated_at: Datetime,
    pub price_per_unit: u64,
    pub locked_nano: u64,
    pub collected_at: Datetime,
}

impl VmContractWithNode {
    pub async fn get_by_uuid(uuid: &str) -> Result<Option<Self>, Error> {
        let contract: Option<Self> =
            DB.query(format!("select * from {VM_CONTRACT}:{uuid} fetch out;")).await?.take(0)?;
        Ok(contract)
    }

    pub async fn list_by_admin(admin: &str) -> Result<Vec<Self>, Error> {
        let mut result = DB
            .query(format!("select * from {VM_CONTRACT} where in = {ACCOUNT}:{admin} fetch out;"))
            .await?;
        let contracts: Vec<Self> = result.take(0)?;
        Ok(contracts)
    }

    pub async fn list_by_operator(operator: &str) -> Result<Vec<Self>, Error> {
        let mut result = DB
            .query(format!(
                "select
                (select * from ->operator->vm_node<-vm_contract fetch out) as contracts
            from {ACCOUNT}:{operator};"
            ))
            .await?;

        #[derive(Deserialize)]
        struct Wrapper {
            contracts: Vec<VmContractWithNode>,
        }

        let c: Option<Wrapper> = result.take(0)?;
        match c {
            Some(c) => Ok(c.contracts),
            None => Ok(Vec::new()),
        }
    }

    /// total hardware units of this VM
    fn total_units(&self) -> u64 {
        // TODO: Optimize this based on price of hardware.
        // I tried, but this can be done better.
        // Storage cost should also be based on tier
        (self.vcpus as u64 * 10)
            + ((self.memory_mb + 256) as u64 / 200)
            + (self.disk_size_gb as u64 / 10)
            + (!self.public_ipv4.is_empty() as u64 * 10)
    }

    /// Returns price per minute in nanoLP
    pub fn price_per_minute(&self) -> u64 {
        self.total_units() * self.price_per_unit
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct AppNode {
    id: RecordId,
    country: String,
    region: String,
    city: String,
    ip: String,
    avail_mem_mb: u32,
    avail_vcpus: u32,
    avail_storage_gbs: u32,
    avail_ports: u32,
    max_ports_per_app: u32,
    price: u64,
    offline_minutes: u64,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct AppContract {
    id: RecordId,
    #[serde(rename = "in")]
    admin: RecordId,
    #[serde(rename = "out")]
    app_node: RecordId,
    state: String,
    app_name: String,
    mapped_ports: Vec<(u32, u32)>,
    host_ipv4: String,
    vcpus: u32,
    memory_mb: u32,
    disk_size_gb: u32,
    created_at: Datetime,
    updated_at: Datetime,
    price_per_unit: u64,
    locked_nano: u64,
    collected_at: Datetime,
    mr_enclave: String,
    package_url: String,
    hratls_pubkey: String,
}

#[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,
    reason: String,
}

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

#[derive(Debug, Serialize, Deserialize)]
pub struct OperatorRelation {
    #[serde(rename = "in")]
    pub account: RecordId,
    #[serde(rename = "out")]
    pub node: RecordId,
}

impl OperatorRelation {
    fn new(account: &str, vm_node: &str) -> Self {
        Self {
            account: RecordId::from(("account", account.to_string())),
            node: RecordId::from(("vm_node", vm_node.to_string())),
        }
    }
}

/// 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() -> Result<Vec<Self>, Error> {
        let mut result = DB
            .query(format!(
                "select *,
                    in as account,
                    <-account.email[0] as email,
                    <-account.escrow[0] as escrow,
                    count(->vm_node) as vm_nodes, 
                    count(->app_node) as app_nodes,
                    (select in from <-account->operator->vm_node<-report).len() +
                        (select in from <-account->operator->app_node<-report).len()
                    as reports
                from operator group by in;"
            ))
            .await?;
        let operators: Vec<Self> = result.take(0)?;
        Ok(operators)
    }
}

// TODO: delete all of these From implementation after migration 0 gets executed

impl From<&old_brain::BrainData> for Vec<VmNode> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut nodes = Vec::new();
        for old_node in old_data.vm_nodes.iter() {
            nodes.push(VmNode {
                id: RecordId::from(("vm_node", old_node.public_key.clone())),
                country: old_node.country.clone(),
                region: old_node.region.clone(),
                city: old_node.city.clone(),
                ip: old_node.ip.clone(),
                avail_mem_mb: old_node.avail_mem_mb,
                avail_vcpus: old_node.avail_vcpus,
                avail_storage_gbs: old_node.avail_storage_gbs,
                avail_ipv4: old_node.avail_ipv4,
                avail_ipv6: old_node.avail_ipv6,
                avail_ports: old_node.avail_ports,
                max_ports_per_vm: old_node.max_ports_per_vm,
                price: old_node.price,
                offline_minutes: old_node.offline_minutes,
            });
        }
        nodes
    }
}

impl From<&old_brain::BrainData> for Vec<VmContract> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut contracts = Vec::new();
        for old_c in old_data.vm_contracts.iter() {
            let mut mapped_ports = Vec::new();
            for port in old_c.exposed_ports.iter() {
                mapped_ports.push((*port, 8080 as u32));
            }
            contracts.push(VmContract {
                id: RecordId::from((VM_CONTRACT, old_c.uuid.replace("-", ""))),
                admin: RecordId::from((ACCOUNT, old_c.admin_pubkey.clone())),
                vm_node: RecordId::from((VM_NODE, old_c.node_pubkey.clone())),
                state: "active".to_string(),
                hostname: old_c.hostname.clone(),
                mapped_ports,
                public_ipv4: old_c.public_ipv4.clone(),
                public_ipv6: old_c.public_ipv6.clone(),
                disk_size_gb: old_c.disk_size_gb,
                vcpus: old_c.vcpus,
                memory_mb: old_c.memory_mb,
                dtrfs_sha: old_c.dtrfs_sha.clone(),
                kernel_sha: old_c.kernel_sha.clone(),
                price_per_unit: old_c.price_per_unit,
                locked_nano: old_c.locked_nano,
                created_at: old_c.created_at.into(),
                updated_at: old_c.updated_at.into(),
                collected_at: old_c.collected_at.into(),
            });
        }
        contracts
    }
}

impl From<&old_brain::BrainData> for Vec<AppNode> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut nodes = Vec::new();
        for old_node in old_data.app_nodes.iter() {
            nodes.push(AppNode {
                id: RecordId::from(("app_node", old_node.node_pubkey.clone())),
                country: old_node.country.clone(),
                region: old_node.region.clone(),
                city: old_node.city.clone(),
                ip: old_node.ip.clone(),
                avail_mem_mb: old_node.avail_mem_mb,
                avail_vcpus: old_node.avail_vcpus,
                avail_storage_gbs: old_node.avail_storage_mb,
                avail_ports: old_node.avail_no_of_port,
                max_ports_per_app: old_node.max_ports_per_app,
                price: old_node.price,
                offline_minutes: old_node.offline_minutes,
            });
        }
        nodes
    }
}

impl From<&old_brain::BrainData> for Vec<AppContract> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut contracts = Vec::new();
        for old_c in old_data.app_contracts.iter() {
            let mut mapped_ports = Vec::new();
            for port in old_c.mapped_ports.clone().into_iter().map(|(b, c)| (b as u32, c as u32)) {
                mapped_ports.push(port);
            }

            let mr_enclave_hex = old_c
                .public_package_mr_enclave
                .clone()
                .unwrap_or_default()
                .iter()
                .map(|byte| format!("{:02X}", byte))
                .collect();

            contracts.push(AppContract {
                id: RecordId::from(("app_contract", old_c.uuid.replace("-", ""))),
                admin: RecordId::from(("account", old_c.admin_pubkey.clone())),
                app_node: RecordId::from(("app_node", old_c.node_pubkey.clone())),
                state: "active".to_string(),
                mapped_ports,
                host_ipv4: old_c.host_ipv4.clone(),
                disk_size_gb: old_c.disk_size_mb * 1024,
                vcpus: old_c.vcpus,
                memory_mb: old_c.memory_mb,
                price_per_unit: old_c.price_per_unit,
                locked_nano: old_c.locked_nano,
                created_at: old_c.created_at.into(),
                updated_at: old_c.updated_at.into(),
                collected_at: old_c.collected_at.into(),
                app_name: old_c.app_name.clone(),
                mr_enclave: mr_enclave_hex,
                package_url: old_c.package_url.clone(),
                hratls_pubkey: old_c.hratls_pubkey.clone(),
            });
        }
        contracts
    }
}

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
    }
}

impl From<&old_brain::BrainData> for Vec<OperatorRelation> {
    fn from(old_data: &old_brain::BrainData) -> Self {
        let mut operator_entries = Vec::new();
        for operator in old_data.operators.clone() {
            for vm_node in operator.1.vm_nodes.iter() {
                operator_entries.push(OperatorRelation::new(&operator.0, vm_node));
            }
            for app_node in operator.1.app_nodes.iter() {
                operator_entries.push(OperatorRelation::new(&operator.0, app_node));
            }
        }
        operator_entries
    }
}