use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StoragePlan {
pub partition: PartitionStrategy,
pub compression: CompressionStrategy,
pub index: IndexStrategy,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "strategy", rename_all = "snake_case")]
pub enum PartitionStrategy {
None,
RangeByTimeDimension {
field: String,
granularity: PartitionGranularity,
#[serde(default, skip_serializing_if = "Option::is_none")]
retention: Option<RetentionPolicy>,
#[serde(default, skip_serializing_if = "Option::is_none")]
look_ahead: Option<u32>,
#[serde(default, skip_serializing_if = "PartitionWindow::is_default")]
window: PartitionWindow,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct PartitionWindow {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub history_days: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub forecast_days: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub partition_budget: Option<u32>,
}
impl PartitionWindow {
#[must_use]
pub fn is_default(&self) -> bool {
self == &Self::default()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PartitionGranularity {
Day,
Week,
Month,
Quarter,
Year,
}
impl PartitionGranularity {
#[must_use]
pub const fn next_coarser(self) -> Option<Self> {
match self {
Self::Day => Some(Self::Week),
Self::Week => Some(Self::Month),
Self::Month => Some(Self::Quarter),
Self::Quarter => Some(Self::Year),
Self::Year => None,
}
}
#[must_use]
pub const fn approx_days(self) -> u32 {
match self {
Self::Day => 1,
Self::Week => 7,
Self::Month => 30,
Self::Quarter => 92,
Self::Year => 365,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RetentionPolicy {
pub keep: u32,
#[serde(default)]
pub action: RetentionAction,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RetentionAction {
#[default]
Drop,
Archive,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum CompressionStrategy {
None,
Page,
Row,
Columnstore,
ColumnstoreArchive,
Lz4,
Zstd {
#[serde(default, skip_serializing_if = "Option::is_none")]
level: Option<u8>,
},
Snappy,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum IndexStrategy {
None,
ClusteredColumnstore,
NonClusteredColumnstore,
ZOrder {
columns: Vec<String>,
},
ClusterBy {
columns: Vec<String>,
},
}
impl StoragePlan {
#[must_use]
pub fn none() -> Self {
Self {
partition: PartitionStrategy::None,
compression: CompressionStrategy::None,
index: IndexStrategy::None,
}
}
#[must_use]
pub fn is_partitioned(&self) -> bool {
!matches!(self.partition, PartitionStrategy::None)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PartitionMaintenance {
pub extend_proc_name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub retain_proc_name: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn storage_plan_none_is_empty() {
let plan = StoragePlan::none();
assert!(!plan.is_partitioned());
assert!(matches!(plan.partition, PartitionStrategy::None));
assert!(matches!(plan.compression, CompressionStrategy::None));
assert!(matches!(plan.index, IndexStrategy::None));
}
#[test]
fn range_partition_round_trips_via_yaml() {
let plan = StoragePlan {
partition: PartitionStrategy::RangeByTimeDimension {
field: "sale_date".into(),
granularity: PartitionGranularity::Day,
retention: Some(RetentionPolicy { keep: 730, action: RetentionAction::Drop }),
look_ahead: Some(7),
window: PartitionWindow::default(),
},
compression: CompressionStrategy::Page,
index: IndexStrategy::ClusteredColumnstore,
};
let yaml = serde_yaml::to_string(&plan).unwrap();
let parsed: StoragePlan = serde_yaml::from_str(&yaml).unwrap();
assert_eq!(plan, parsed);
assert!(plan.is_partitioned());
}
#[test]
fn partition_window_default_is_skipped_in_yaml() {
let plan = StoragePlan {
partition: PartitionStrategy::RangeByTimeDimension {
field: "sale_date".into(),
granularity: PartitionGranularity::Day,
retention: None,
look_ahead: None,
window: PartitionWindow::default(),
},
compression: CompressionStrategy::None,
index: IndexStrategy::None,
};
let yaml = serde_yaml::to_string(&plan).unwrap();
assert!(!yaml.contains("window"), "default window should not serialize: {yaml}");
}
#[test]
fn partition_window_round_trips_via_yaml() {
let window = PartitionWindow {
history_days: Some(1825),
forecast_days: Some(365),
partition_budget: Some(2000),
};
let plan = StoragePlan {
partition: PartitionStrategy::RangeByTimeDimension {
field: "sale_date".into(),
granularity: PartitionGranularity::Day,
retention: None,
look_ahead: None,
window: window.clone(),
},
compression: CompressionStrategy::None,
index: IndexStrategy::None,
};
let yaml = serde_yaml::to_string(&plan).unwrap();
let parsed: StoragePlan = serde_yaml::from_str(&yaml).unwrap();
assert_eq!(plan, parsed);
let PartitionStrategy::RangeByTimeDimension { window: w, .. } = parsed.partition else {
unreachable!()
};
assert_eq!(w, window);
}
#[test]
fn next_coarser_walks_to_year_then_stops() {
assert_eq!(PartitionGranularity::Day.next_coarser(), Some(PartitionGranularity::Week));
assert_eq!(PartitionGranularity::Week.next_coarser(), Some(PartitionGranularity::Month));
assert_eq!(PartitionGranularity::Month.next_coarser(), Some(PartitionGranularity::Quarter));
assert_eq!(PartitionGranularity::Quarter.next_coarser(), Some(PartitionGranularity::Year));
assert_eq!(PartitionGranularity::Year.next_coarser(), None);
}
#[test]
fn zstd_with_level_serializes() {
let c = CompressionStrategy::Zstd { level: Some(3) };
let json = serde_json::to_string(&c).unwrap();
assert!(json.contains("\"zstd\""));
assert!(json.contains('3'));
}
#[test]
fn retention_action_defaults_to_drop() {
let yaml = "keep: 90\n";
let policy: RetentionPolicy = serde_yaml::from_str(yaml).unwrap();
assert_eq!(policy.keep, 90);
assert_eq!(policy.action, RetentionAction::Drop);
}
#[test]
fn z_order_round_trips_via_json() {
let idx = IndexStrategy::ZOrder { columns: vec!["sale_date".into(), "store_id".into()] };
let json = serde_json::to_string(&idx).unwrap();
let parsed: IndexStrategy = serde_json::from_str(&json).unwrap();
assert_eq!(idx, parsed);
}
#[test]
fn cluster_by_round_trips_via_json() {
let idx = IndexStrategy::ClusterBy { columns: vec!["created_at".into()] };
let json = serde_json::to_string(&idx).unwrap();
let parsed: IndexStrategy = serde_json::from_str(&json).unwrap();
assert_eq!(idx, parsed);
}
#[test]
fn partition_strategy_none_serializes_with_tag() {
let p = PartitionStrategy::None;
let json = serde_json::to_string(&p).unwrap();
assert!(json.contains("\"none\""));
}
#[test]
fn partition_maintenance_round_trips() {
let m = PartitionMaintenance {
extend_proc_name: "core_sales_partition_extend".into(),
retain_proc_name: Some("core_sales_partition_retain".into()),
};
let json = serde_json::to_string(&m).unwrap();
let parsed: PartitionMaintenance = serde_json::from_str(&json).unwrap();
assert_eq!(m, parsed);
let no_retain = PartitionMaintenance {
extend_proc_name: "core_orders_partition_extend".into(),
retain_proc_name: None,
};
let json = serde_json::to_string(&no_retain).unwrap();
assert!(!json.contains("retain_proc_name"));
}
}