use anyhow::Error;
use surrealdb::{
engine::local::{Db, RocksDb},
Surreal,
};
use tracing::{debug, error};
use super::super::order::{Amount, Crypto, Order};
const DB_NAMESPACE: &str = "scuriolus";
const TABLE_ORDRES: &str = "orders";
#[derive(Debug)]
pub struct Data {
db: Surreal<Db>,
}
impl Data {
pub async fn load_data(name: String, address: String) -> Result<Self, Error> {
let db = Surreal::new::<RocksDb>(address).await?;
db.use_ns(DB_NAMESPACE).use_db(name).await?;
Ok(Data { db })
}
pub async fn add_order(&self, order: Order) -> Result<(), Error> {
let previous: Option<Order> = self.db.select((TABLE_ORDRES, order.id())).await?;
if previous.is_some() {
error!("Add order: {} already exists", order.id());
debug!("Previous: {:#?}", previous.unwrap());
debug!("Challenger: {:#?}", order);
return Err(Error::msg("id already occupied"));
}
let _: Option<Order> = self
.db
.create((TABLE_ORDRES, order.id()))
.content(order)
.await?;
Ok(())
}
pub async fn get_order(&self, id: &str) -> Option<Order> {
self.db.select((TABLE_ORDRES, id)).await.unwrap()
}
pub async fn update_order(&self, order: Order) -> Result<(), Error> {
let previous: Option<Order> = self.db.select((TABLE_ORDRES, order.id())).await?;
if previous.is_none() {
error!("Update order: {} not found", order.id());
debug!("Order: {:#?}", order);
return Err(Error::msg("order not found"));
}
let _updated: Option<Order> = self
.db
.upsert((TABLE_ORDRES, order.id()))
.content(order)
.await?;
Ok(())
}
pub async fn opened_orders(&self) -> Vec<Order> {
let query = "SELECT * from orders WHERE status IN ['Open', 'PartiallyFilled']";
let mut res = self.db.query(query).await.unwrap();
let orders: Vec<Order> = res.take(0).expect("Corrupted DataBase");
orders
}
#[cfg(test)]
pub async fn clear(&self) {
let _: Result<Vec<Order>, surrealdb::Error> = self.db.delete(TABLE_ORDRES).await;
let _: Result<Vec<Order>, surrealdb::Error> = self.db.delete(TABLE_ASSETS).await;
}
}
const TABLE_ASSETS: &str = "main_assets";
#[derive(
serde::Serialize, serde::Deserialize, Debug, derive_getters::Getters, Clone, Copy, PartialEq,
)]
pub struct Asset {
crypto: Crypto,
pub free: Amount,
pub locked: Amount,
}
impl Asset {
pub fn new(crypto: Crypto) -> Self {
Self {
crypto,
free: Amount::from(0),
locked: Amount::from(0),
}
}
}
#[allow(dead_code)]
pub trait DataAsset {
async fn get_asset(&self, crypto: &Crypto) -> Result<Asset, Error>;
async fn update_asset(&self, asset: Asset) -> Result<(), Error>;
}
impl DataAsset for Data {
async fn get_asset(&self, crypto: &Crypto) -> Result<Asset, Error> {
let o_asset: Option<Asset> = self.db.select((TABLE_ASSETS, crypto.to_string())).await?;
if let Some(asset) = o_asset {
Ok(asset)
} else {
tracing::debug!("Asset {} not found, creating it", crypto);
Ok(Asset::new(*crypto))
}
}
async fn update_asset(&self, asset: Asset) -> Result<(), Error> {
let _: Option<Asset> = self
.db
.upsert((TABLE_ASSETS, asset.crypto().to_string()))
.content(asset)
.await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use serial_test::serial;
use super::super::super::{
market::mexc_base::{mexc_enums, mexc_enums::OrderType},
order::{Amount, Crypto, OrderStatus, Quantity},
};
use super::*;
const DB_ADDRESS: &str = "db/test/";
#[tokio::test]
#[serial]
async fn orders_basics() {
let data = Data::load_data("test".to_string(), DB_ADDRESS.to_string())
.await
.unwrap();
data.clear().await;
let mut order = Order::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Limit,
Quantity::Asset(Amount::ONE),
Some(Amount::from(10000)),
)
.unwrap();
let mut order2 = Order::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Market,
Quantity::Asset(Amount::ONE),
None,
)
.unwrap();
data.add_order(order.clone()).await.unwrap();
data.add_order(order2.clone()).await.unwrap();
let recup = data.get_order(order.id()).await.unwrap();
assert_eq!(recup, order);
order.set_status(OrderStatus::Tested);
data.update_order(order.clone()).await.unwrap();
let recup = data.get_order(order.id()).await.unwrap();
assert_eq!(recup, order);
order2.set_status(OrderStatus::Tested);
order2.set_price(Amount::from(20000)).unwrap();
data.update_order(order2.clone()).await.unwrap();
let recup2 = data.get_order(order2.id()).await.unwrap();
assert_eq!(recup2, order2);
}
#[tokio::test]
#[serial]
async fn orders_get_opens() {
let data = Data::load_data("test".to_string(), DB_ADDRESS.to_string())
.await
.unwrap();
data.clear().await;
let order_open = Order::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Limit,
Quantity::Asset(Amount::ONE),
Some(Amount::from(10000)),
)
.unwrap();
let mut order_part = Order::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Limit,
Quantity::Asset(Amount::ONE),
Some(Amount::from(10000)),
)
.unwrap();
order_part.set_status(OrderStatus::PartiallyFilled);
let mut order_executed = Order::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Market,
Quantity::Asset(Amount::ONE),
None,
)
.unwrap();
order_executed.set_status(OrderStatus::Filled);
tracing::debug!("Order open: {}", order_open);
data.add_order(order_open).await.unwrap();
data.add_order(order_part).await.unwrap();
data.add_order(order_executed).await.unwrap();
let open_orders = data.opened_orders().await;
assert_eq!(open_orders.len(), 2);
}
#[tokio::test]
#[serial]
async fn assets() {
let data = Data::load_data("test".to_string(), DB_ADDRESS.to_string())
.await
.unwrap();
data.clear().await;
let crypto = Crypto::BTC;
let mut asset = Asset::new(crypto);
asset.free = Amount::from(3);
asset.locked = Amount::from(4);
data.update_asset(asset).await.unwrap();
let recup = data.get_asset(&crypto).await.unwrap();
assert_eq!(recup, asset);
asset.free = Amount::from(42);
data.update_asset(asset).await.unwrap();
let recup = data.get_asset(&crypto).await.unwrap();
assert_eq!(recup, asset);
}
}