use polars::prelude::*;
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Error, Debug)]
pub enum OpenSkyError {
#[error("Configuration error: {0}")]
Config(String),
#[error("Authentication failed: {0}")]
Auth(String),
#[error("HTTP request failed: {0}")]
Http(#[from] reqwest::Error),
#[error("Query execution failed: {0}")]
Query(String),
#[error("Query was cancelled")]
Cancelled,
#[error("Invalid parameter: {0}")]
InvalidParam(String),
#[error("Data conversion error: {0}")]
DataConversion(String),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("JSON parsing error: {0}")]
Json(#[from] serde_json::Error),
}
pub type Result<T> = std::result::Result<T, OpenSkyError>;
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Bounds {
pub west: f64,
pub south: f64,
pub east: f64,
pub north: f64,
}
impl Bounds {
pub fn new(west: f64, south: f64, east: f64, north: f64) -> Self {
Self { west, south, east, north }
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct QueryParams {
pub icao24: Option<String>,
pub start: Option<String>,
pub stop: Option<String>,
pub callsign: Option<String>,
pub bounds: Option<Bounds>,
pub departure_airport: Option<String>,
pub arrival_airport: Option<String>,
pub airport: Option<String>,
pub time_buffer: Option<String>,
pub limit: Option<u32>,
}
impl QueryParams {
pub fn new() -> Self {
Self::default()
}
pub fn icao24(mut self, icao24: impl Into<String>) -> Self {
self.icao24 = Some(icao24.into());
self
}
pub fn time_range(mut self, start: impl Into<String>, stop: impl Into<String>) -> Self {
self.start = Some(start.into());
self.stop = Some(stop.into());
self
}
pub fn departure(mut self, airport: impl Into<String>) -> Self {
self.departure_airport = Some(airport.into());
self
}
pub fn arrival(mut self, airport: impl Into<String>) -> Self {
self.arrival_airport = Some(airport.into());
self
}
pub fn limit(mut self, limit: u32) -> Self {
self.limit = Some(limit);
self
}
pub fn bounds(mut self, west: f64, south: f64, east: f64, north: f64) -> Self {
self.bounds = Some(Bounds::new(west, south, east, north));
self
}
pub fn is_empty(&self) -> bool {
self.icao24.is_none()
&& self.start.is_none()
&& self.stop.is_none()
&& self.callsign.is_none()
&& self.bounds.is_none()
&& self.departure_airport.is_none()
&& self.arrival_airport.is_none()
&& self.airport.is_none()
}
}
pub const FLIGHT_COLUMNS: &[&str] = &[
"time",
"icao24",
"lat",
"lon",
"velocity",
"heading",
"vertrate",
"callsign",
"onground",
"squawk",
"baroaltitude",
"geoaltitude",
"hour",
];
pub const FLIGHTLIST_COLUMNS: &[&str] = &[
"icao24",
"callsign",
"firstseen",
"lastseen",
"estdepartureairport",
"estarrivalairport",
"day",
];
pub const RAWDATA_COLUMNS: &[&str] = &[
"mintime",
"rawmsg",
"icao24",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum RawTable {
#[default]
RollcallReplies,
Acas,
AllcallReplies,
Identification,
OperationalStatus,
Position,
Velocity,
}
impl RawTable {
pub fn table_name(&self) -> &'static str {
match self {
RawTable::RollcallReplies => "minio.osky.rollcall_replies_data4",
RawTable::Acas => "minio.osky.acas_data4",
RawTable::AllcallReplies => "minio.osky.allcall_replies_data4",
RawTable::Identification => "minio.osky.identification_data4",
RawTable::OperationalStatus => "minio.osky.operational_status_data4",
RawTable::Position => "minio.osky.position_data4",
RawTable::Velocity => "minio.osky.velocity_data4",
}
}
}
#[derive(Debug, Clone)]
pub struct FlightData {
df: DataFrame,
}
impl FlightData {
pub fn new(df: DataFrame) -> Self {
Self { df }
}
pub fn dataframe(&self) -> &DataFrame {
&self.df
}
pub fn dataframe_mut(&mut self) -> &mut DataFrame {
&mut self.df
}
pub fn into_dataframe(self) -> DataFrame {
self.df
}
pub fn len(&self) -> usize {
self.df.height()
}
pub fn is_empty(&self) -> bool {
self.df.height() == 0
}
pub fn columns(&self) -> Vec<String> {
self.df.get_column_names().iter().map(|s| s.to_string()).collect()
}
pub fn to_csv(&self, path: &str) -> Result<()> {
let mut file = std::fs::File::create(path)?;
CsvWriter::new(&mut file)
.finish(&mut self.df.clone())
.map_err(|e| OpenSkyError::DataConversion(e.to_string()))?;
Ok(())
}
pub fn to_parquet(&self, path: impl AsRef<std::path::Path>) -> Result<()> {
let mut file = std::fs::File::create(path)?;
ParquetWriter::new(&mut file)
.finish(&mut self.df.clone())
.map_err(|e| OpenSkyError::DataConversion(e.to_string()))?;
Ok(())
}
pub fn from_parquet(path: impl AsRef<std::path::Path>) -> Result<Self> {
let file = std::fs::File::open(path)?;
let df = ParquetReader::new(file)
.finish()
.map_err(|e| OpenSkyError::DataConversion(e.to_string()))?;
Ok(Self { df })
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_query_params_builder() {
let params = QueryParams::new()
.icao24("485a32")
.time_range("2025-01-01 00:00:00", "2025-01-01 23:59:59")
.departure("EHAM")
.arrival("EGLL");
assert_eq!(params.icao24, Some("485a32".to_string()));
assert_eq!(params.departure_airport, Some("EHAM".to_string()));
assert!(!params.is_empty());
}
#[test]
fn test_query_params_empty() {
let params = QueryParams::new();
assert!(params.is_empty());
}
}