use crate::utils::current_timestamp_ms;
use crate::{
errors::PersistError,
exchanges::Exchange,
orderbooks::{OrderbookDelta, io},
};
use serde::{Deserialize, Serialize};
use std::path::Path;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutputFormat {
Csv,
Json,
Parquet,
}
impl OutputFormat {
pub fn from_path(path: &Path) -> Option<Self> {
match path.extension()?.to_str()? {
"csv" => Some(Self::Csv),
"json" => Some(Self::Json),
"parquet" | "pq" => Some(Self::Parquet),
_ => None,
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct PriceLevelRecord {
pub timestamp_ms: u64,
pub symbol: String,
pub side: String, pub level: usize, pub price: String, pub size: String,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct OrderbookSnapshot {
pub timestamp_ms: u64,
pub symbol: String,
pub update_id: u64,
pub sequence: u64,
pub delta_count: u64,
pub bid_depth: usize,
pub ask_depth: usize,
pub mid_price: Option<String>,
pub spread: Option<String>,
pub spread_bps: Option<String>,
pub volume_imbalance: Option<String>,
pub total_bid_volume: String,
pub total_ask_volume: String,
pub bids: Vec<[String; 2]>, pub asks: Vec<[String; 2]>,
}
pub fn save_orderbook_state(
ob: &OrderbookDelta,
path: impl AsRef<Path>,
format: Option<OutputFormat>,
) -> Result<(), PersistError> {
let path = path.as_ref();
let format = format
.or_else(|| OutputFormat::from_path(path))
.ok_or_else(|| {
PersistError::UnsupportedFormat(
path.extension()
.and_then(|e| e.to_str())
.unwrap_or("unknown")
.to_string(),
)
})?;
match format {
OutputFormat::Csv => io::write_csv(ob, path),
OutputFormat::Json => io::write_json(ob, path),
OutputFormat::Parquet => io::write_ob_delta_parquet(ob, path),
}
}
pub fn save_orderbook_timestamped(
ob: &OrderbookDelta,
ex: &Exchange,
dir: impl AsRef<Path>,
format: OutputFormat,
) -> Result<std::path::PathBuf, PersistError> {
let timestamp = current_timestamp_ms();
let ext = match format {
OutputFormat::Csv => "csv",
OutputFormat::Json => "json",
OutputFormat::Parquet => "parquet",
};
let exchange_name: String = match ex {
Exchange::Bybit => "bybit".to_string(),
Exchange::Coinbase => "coinbase".to_string(),
Exchange::Kraken => "kraken".to_string(),
Exchange::Binance => "binance".to_string(),
};
let filename = format!(
"{}_{}_{:?}.{}",
exchange_name.to_lowercase(),
ob.symbol().to_lowercase(),
timestamp,
ext
);
let path = dir.as_ref().join(filename);
save_orderbook_state(ob, &path, Some(format))?;
Ok(path)
}