use rust_decimal::Decimal;
use surrealdb::{
engine::local::{Db, RocksDb},
Surreal,
};
use tracing::{debug, error};
use super::core_error::{CoreError, CoreResult};
use crate::data::{Crypto, Order, SpecificOrderDetails};
const DB_NAMESPACE: &str = "scuriolus";
const DB_NAME: &str = "main";
#[cfg(not(test))]
const DB_ADDRESS_BASE: &str = "db/storage";
#[cfg(test)]
const DB_ADDRESS_BASE: &str = "db/test-storage";
const TABLE_ORDRES: &str = "orders";
#[derive(Debug)]
pub struct Storage {
name: String,
db: Surreal<Db>,
}
impl Storage {
pub async fn new(name: &str) -> CoreResult<Self> {
let db = Surreal::new::<RocksDb>(format!("{}/{}", DB_ADDRESS_BASE, name)).await?;
db.use_ns(DB_NAMESPACE).use_db(DB_NAME).await?;
Ok(Storage {
name: name.to_string(),
db,
})
}
pub fn get_name(&self) -> &str {
&self.name
}
pub async fn add_order<D: SpecificOrderDetails>(&self, order: Order<D>) -> CoreResult<()> {
tracing::debug!("Adding order {} to storage", order.id());
let previous: Option<Order<D>> = 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(CoreError::param_error("id already occupied"));
}
let _: Option<Order<D>> = self
.db
.create((TABLE_ORDRES, order.id()))
.content(order)
.await?;
Ok(())
}
pub async fn get_order<D: SpecificOrderDetails>(&self, id: &str) -> Option<Order<D>> {
self.db.select((TABLE_ORDRES, id)).await.unwrap()
}
pub async fn update_order<D: SpecificOrderDetails>(&self, order: Order<D>) -> CoreResult<()> {
let previous: Option<Order<D>> = self.db.select((TABLE_ORDRES, order.id())).await?;
if previous.is_none() {
error!("Update order: {} not found", order.id());
debug!("Order: {:#?}", order);
return Err(CoreError::param_error("order not found"));
}
let _updated: Option<Order<D>> = self
.db
.upsert((TABLE_ORDRES, order.id()))
.content(order)
.await?;
Ok(())
}
pub async fn opened_orders<D: SpecificOrderDetails>(&self) -> Vec<Order<D>> {
let query = "SELECT * from orders WHERE status IN ['Open', 'PartiallyFilled']";
let mut res = self.db.query(query).await.unwrap();
let orders: Vec<Order<D>> = res.take(0).expect("Corrupted DataBase");
orders
}
#[cfg(test)]
pub async fn clear<D: SpecificOrderDetails>(&self) -> CoreResult<()> {
let _: Vec<Order<D>> = self.db.delete(TABLE_ORDRES).await?;
let _: Vec<Asset> = self.db.delete(TABLE_ASSETS).await?;
Ok(())
}
}
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: Decimal,
pub locked: Decimal,
}
impl Asset {
pub fn new(crypto: Crypto) -> Self {
Self {
crypto,
free: Decimal::from(0),
locked: Decimal::from(0),
}
}
}
#[allow(dead_code)]
pub trait DataAsset {
async fn get_asset(&self, crypto: &Crypto) -> CoreResult<Asset>;
async fn update_asset(&self, asset: Asset) -> CoreResult<()>;
}
impl DataAsset for Storage {
async fn get_asset(&self, crypto: &Crypto) -> CoreResult<Asset> {
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) -> CoreResult<()> {
let _: Option<Asset> = self
.db
.upsert((TABLE_ASSETS, asset.crypto().to_string()))
.content(asset)
.await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::data::EmptySpecificOrderDetails;
use super::super::super::{
data::{Crypto, OrderStatus, Quantity},
market::{mexc_enums, mexc_enums::OrderType},
};
use super::*;
#[tokio::test]
async fn orders_basics() {
let data = Storage::new("data-order_basic").await.unwrap();
data.clear::<EmptySpecificOrderDetails>().await.unwrap();
let mut order = Order::<EmptySpecificOrderDetails>::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Limit,
Quantity::Asset(Decimal::ONE),
Some(Decimal::from(10000)),
)
.unwrap();
let mut order2 = Order::<EmptySpecificOrderDetails>::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Market,
Quantity::Asset(Decimal::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(Decimal::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]
async fn orders_get_opens() {
let data = Storage::new("data-orders_get_opens").await.unwrap();
data.clear::<EmptySpecificOrderDetails>().await.unwrap();
let order_open = Order::<EmptySpecificOrderDetails>::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Limit,
Quantity::Asset(Decimal::ONE),
Some(Decimal::from(10000)),
)
.unwrap();
let mut order_part = Order::<EmptySpecificOrderDetails>::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Limit,
Quantity::Asset(Decimal::ONE),
Some(Decimal::from(10000)),
)
.unwrap();
order_part.set_status(OrderStatus::PartiallyFilled);
let mut order_executed = Order::<EmptySpecificOrderDetails>::new(
Crypto::BTC,
Crypto::USDT,
mexc_enums::OrderSide::Buy,
OrderType::Market,
Quantity::Asset(Decimal::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::<EmptySpecificOrderDetails>().await;
assert_eq!(open_orders.len(), 2);
}
#[tokio::test]
async fn assets() {
let data = Storage::new("data-assets").await.unwrap();
data.clear::<EmptySpecificOrderDetails>().await.unwrap();
let crypto = Crypto::BTC;
let mut asset = Asset::new(crypto);
asset.free = Decimal::from(3);
asset.locked = Decimal::from(4);
data.update_asset(asset).await.unwrap();
let recup = data.get_asset(&crypto).await.unwrap();
assert_eq!(recup, asset);
asset.free = Decimal::from(42);
data.update_asset(asset).await.unwrap();
let recup = data.get_asset(&crypto).await.unwrap();
assert_eq!(recup, asset);
}
}