use std::{cmp::Ordering, collections::HashSet, fmt, num::NonZeroU64, str::FromStr};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use snafu::prelude::*;
#[derive(Debug, Snafu, PartialEq, Eq)]
#[non_exhaustive]
pub enum ParseTimeIndexGranularityError {
#[snafu(display("time index granularity is empty"))]
Empty,
#[snafu(display("time index granularity '{spec}' is missing a numeric value"))]
MissingNumber {
spec: String,
},
#[snafu(display(
"time index granularity '{spec}' is missing a unit suffix (expected s|m|h|d)"
))]
MissingUnit {
spec: String,
},
#[snafu(display("invalid index granularity value in '{spec}': {source}"))]
InvalidNumber {
spec: String,
source: std::num::ParseIntError,
},
#[snafu(display("index granularity value must be > 0 (got {value}) in '{spec}'"))]
NonPositive {
spec: String,
value: u64,
},
#[snafu(display("index granularity value too large for u32 (got {value}) in '{spec}'"))]
TooLarge {
spec: String,
value: u64,
},
#[snafu(display(
"unknown time index granularity unit '{unit}' in '{spec}' (expected s|m|h|d)"
))]
UnknownUnit {
spec: String,
unit: String,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum TimeIndexGranularity {
Seconds(u32),
Minutes(u32),
Hours(u32),
Days(u32),
}
impl FromStr for TimeIndexGranularity {
type Err = ParseTimeIndexGranularityError;
fn from_str(input: &str) -> Result<Self, Self::Err> {
let spec = input.trim();
if spec.is_empty() {
return Err(ParseTimeIndexGranularityError::Empty);
}
let unit_start = spec
.char_indices()
.find(|(_, c)| c.is_ascii_alphabetic())
.map(|(i, _)| i);
let Some(unit_start) = unit_start else {
return Err(ParseTimeIndexGranularityError::MissingUnit {
spec: spec.to_string(),
});
};
if unit_start == 0 {
return Err(ParseTimeIndexGranularityError::MissingNumber {
spec: spec.to_string(),
});
}
let (num_str, unit_str) = spec.split_at(unit_start);
let num_str = num_str.trim();
let unit_str = unit_str.trim();
if unit_str.is_empty() {
return Err(ParseTimeIndexGranularityError::MissingUnit {
spec: spec.to_string(),
});
}
let value: u64 =
num_str
.parse()
.map_err(|source| ParseTimeIndexGranularityError::InvalidNumber {
spec: spec.to_string(),
source,
})?;
if value == 0 {
return Err(ParseTimeIndexGranularityError::NonPositive {
spec: spec.to_string(),
value,
});
}
if value > u32::MAX as u64 {
return Err(ParseTimeIndexGranularityError::TooLarge {
spec: spec.to_string(),
value,
});
}
let v = value as u32;
let unit = unit_str.to_ascii_lowercase();
match unit.as_str() {
"s" | "sec" | "secs" | "second" | "seconds" => Ok(TimeIndexGranularity::Seconds(v)),
"m" | "min" | "mins" | "minute" | "minutes" => Ok(TimeIndexGranularity::Minutes(v)),
"h" | "hr" | "hrs" | "hour" | "hours" => Ok(TimeIndexGranularity::Hours(v)),
"d" | "day" | "days" => Ok(TimeIndexGranularity::Days(v)),
_ => Err(ParseTimeIndexGranularityError::UnknownUnit {
spec: spec.to_string(),
unit: unit_str.to_string(),
}),
}
}
}
impl TimeIndexGranularity {
pub fn parse(spec: &str) -> Result<Self, ParseTimeIndexGranularityError> {
spec.parse()
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct IndexSpec {
pub column: String,
#[serde(default)]
pub entity_columns: Vec<String>,
pub kind: IndexKind,
}
impl IndexSpec {
pub fn validate(&self) -> Result<(), IndexSpecError> {
if self.column.is_empty() {
return Err(IndexSpecError::EmptyColumn);
}
let mut seen = HashSet::with_capacity(self.entity_columns.len());
for (position, column) in self.entity_columns.iter().enumerate() {
if column.is_empty() {
return Err(IndexSpecError::EmptyEntityColumn { position });
}
if column == &self.column {
return Err(IndexSpecError::EntityColumnMatchesIndex {
column: column.clone(),
});
}
if !seen.insert(column) {
return Err(IndexSpecError::DuplicateEntityColumn {
column: column.clone(),
});
}
}
self.kind.validate()
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum IndexKind {
Timestamp {
index_granularity: TimeIndexGranularity,
#[serde(default, skip_serializing_if = "Option::is_none")]
timezone: Option<String>,
},
Int64 {
index_granularity: NonZeroU64,
},
#[serde(rename = "uint64")]
UInt64 {
index_granularity: NonZeroU64,
},
}
impl IndexKind {
pub fn name(&self) -> &'static str {
match self {
Self::Timestamp { .. } => "timestamp",
Self::Int64 { .. } => "int64",
Self::UInt64 { .. } => "uint64",
}
}
pub fn validate(&self) -> Result<(), IndexSpecError> {
if let Self::Timestamp {
index_granularity, ..
} = self
{
let width = match index_granularity {
TimeIndexGranularity::Seconds(width)
| TimeIndexGranularity::Minutes(width)
| TimeIndexGranularity::Hours(width)
| TimeIndexGranularity::Days(width) => *width,
};
if width == 0 {
return Err(IndexSpecError::ZeroTimeIndexGranularity);
}
}
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(
tag = "type",
content = "value",
rename_all = "snake_case",
deny_unknown_fields
)]
pub enum IndexValue {
Timestamp(DateTime<Utc>),
Int64(i64),
UInt64(u64),
}
impl IndexValue {
pub fn kind_name(&self) -> &'static str {
match self {
Self::Timestamp(_) => "timestamp",
Self::Int64(_) => "int64",
Self::UInt64(_) => "uint64",
}
}
pub fn compare(&self, other: &Self) -> Result<Ordering, IndexValueError> {
match (self, other) {
(Self::Timestamp(left), Self::Timestamp(right)) => Ok(left.cmp(right)),
(Self::Int64(left), Self::Int64(right)) => Ok(left.cmp(right)),
(Self::UInt64(left), Self::UInt64(right)) => Ok(left.cmp(right)),
_ => Err(IndexValueError::DomainMismatch {
left: self.kind_name(),
right: other.kind_name(),
}),
}
}
pub fn validate_kind(&self, kind: &IndexKind) -> Result<(), IndexValueError> {
if self.kind_name() == kind.name() {
Ok(())
} else {
Err(IndexValueError::KindMismatch {
expected: kind.name(),
actual: self.kind_name(),
})
}
}
}
impl fmt::Display for IndexValue {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Timestamp(value) => write!(f, "timestamp({value})"),
Self::Int64(value) => write!(f, "int64({value})"),
Self::UInt64(value) => write!(f, "uint64({value})"),
}
}
}
impl From<DateTime<Utc>> for IndexValue {
fn from(value: DateTime<Utc>) -> Self {
Self::Timestamp(value)
}
}
impl From<i64> for IndexValue {
fn from(value: i64) -> Self {
Self::Int64(value)
}
}
impl From<u64> for IndexValue {
fn from(value: u64) -> Self {
Self::UInt64(value)
}
}
pub fn validate_index_range(
kind: &IndexKind,
start: &IndexValue,
end: &IndexValue,
) -> Result<(), IndexValueError> {
start.validate_kind(kind)?;
end.validate_kind(kind)?;
if start.compare(end)? != Ordering::Less {
return Err(IndexValueError::InvalidRange {
start: start.clone(),
end: end.clone(),
});
}
Ok(())
}
#[derive(Debug, Snafu, PartialEq, Eq)]
#[non_exhaustive]
pub enum IndexSpecError {
#[snafu(display("ordered index column is empty"))]
EmptyColumn,
#[snafu(display("entity column at position {position} is empty"))]
EmptyEntityColumn {
position: usize,
},
#[snafu(display("duplicate entity column: {column}"))]
DuplicateEntityColumn {
column: String,
},
#[snafu(display("entity column cannot also be the ordered index column: {column}"))]
EntityColumnMatchesIndex {
column: String,
},
#[snafu(display("timestamp index granularity must be nonzero"))]
ZeroTimeIndexGranularity,
}
#[derive(Debug, Snafu, PartialEq, Eq)]
#[non_exhaustive]
pub enum IndexValueError {
#[snafu(display("ordered index domain mismatch: left={left}, right={right}"))]
DomainMismatch {
left: &'static str,
right: &'static str,
},
#[snafu(display("ordered index kind mismatch: expected {expected}, found {actual}"))]
KindMismatch {
expected: &'static str,
actual: &'static str,
},
#[snafu(display(
"invalid ordered index range: start={start}, end={end} (expected start < end)"
))]
InvalidRange {
start: IndexValue,
end: IndexValue,
},
#[snafu(display("invalid ordered index bounds: min={min}, max={max} (expected min <= max)"))]
InvalidBounds {
min: IndexValue,
max: IndexValue,
},
}
#[cfg(test)]
mod tests {
use chrono::TimeZone;
use super::*;
fn sample_time_index_spec() -> IndexSpec {
IndexSpec {
column: "ts".to_string(),
entity_columns: vec!["symbol".to_string()],
kind: IndexKind::Timestamp {
index_granularity: TimeIndexGranularity::Minutes(1),
timezone: None,
},
}
}
#[test]
fn index_spec_json_roundtrips_all_domains() {
let cases = [
(
sample_time_index_spec(),
serde_json::json!({
"column": "ts",
"entity_columns": ["symbol"],
"kind": {
"type": "timestamp",
"index_granularity": {"Minutes": 1}
}
}),
),
(
IndexSpec {
column: "sequence".to_string(),
entity_columns: Vec::new(),
kind: IndexKind::Int64 {
index_granularity: NonZeroU64::new(u64::MAX).unwrap(),
},
},
serde_json::json!({
"column": "sequence",
"entity_columns": [],
"kind": {
"type": "int64",
"index_granularity": u64::MAX
}
}),
),
(
IndexSpec {
column: "offset".to_string(),
entity_columns: vec!["source".to_string()],
kind: IndexKind::UInt64 {
index_granularity: NonZeroU64::new(7).unwrap(),
},
},
serde_json::json!({
"column": "offset",
"entity_columns": ["source"],
"kind": {
"type": "uint64",
"index_granularity": 7
}
}),
),
];
for (spec, expected_json) in cases {
let json = serde_json::to_value(&spec).unwrap();
assert_eq!(json, expected_json);
let restored: IndexSpec = serde_json::from_value(json).unwrap();
assert_eq!(restored, spec);
}
}
#[test]
fn index_spec_json_rejects_impossible_field_combinations() {
let timestamp_with_integer_granularity = r#"{
"column":"ts","kind":{"type":"timestamp","index_granularity":1}
}"#;
let integer_with_time_granularity_object = r#"{
"column":"id","kind":{"type":"int64","index_granularity":{"Seconds":1}}
}"#;
let zero_integer_granularity =
r#"{"column":"id","kind":{"type":"uint64","index_granularity":0}}"#;
assert!(serde_json::from_str::<IndexSpec>(timestamp_with_integer_granularity).is_err());
assert!(serde_json::from_str::<IndexSpec>(integer_with_time_granularity_object).is_err());
assert!(serde_json::from_str::<IndexSpec>(zero_integer_granularity).is_err());
}
#[test]
fn index_spec_validation_rejects_invalid_structure_and_time_granularity() {
let mut spec = sample_time_index_spec();
spec.column.clear();
assert_eq!(spec.validate(), Err(IndexSpecError::EmptyColumn));
let mut spec = sample_time_index_spec();
spec.entity_columns.push("symbol".to_string());
assert!(matches!(
spec.validate(),
Err(IndexSpecError::DuplicateEntityColumn { .. })
));
let mut spec = sample_time_index_spec();
spec.entity_columns = vec![spec.column.clone()];
assert_eq!(
spec.validate(),
Err(IndexSpecError::EntityColumnMatchesIndex {
column: "ts".to_string(),
})
);
let mut spec = sample_time_index_spec();
spec.kind = IndexKind::Timestamp {
index_granularity: TimeIndexGranularity::Seconds(0),
timezone: None,
};
assert_eq!(
spec.validate(),
Err(IndexSpecError::ZeroTimeIndexGranularity)
);
}
#[test]
fn index_value_roundtrips_and_compares_integer_extremes() {
let timestamp = Utc.timestamp_opt(1, 987_654_321).single().unwrap();
let values = [
IndexValue::Timestamp(timestamp),
IndexValue::Int64(i64::MIN),
IndexValue::Int64(i64::MAX),
IndexValue::UInt64(0),
IndexValue::UInt64(u64::MAX),
];
for value in values {
let json = serde_json::to_string(&value).unwrap();
assert_eq!(serde_json::from_str::<IndexValue>(&json).unwrap(), value);
assert_eq!(value.compare(&value).unwrap(), Ordering::Equal);
}
assert_eq!(
IndexValue::Int64(i64::MIN)
.compare(&IndexValue::Int64(i64::MAX))
.unwrap(),
Ordering::Less
);
assert_eq!(
IndexValue::UInt64(u64::MAX)
.compare(&IndexValue::UInt64(0))
.unwrap(),
Ordering::Greater
);
}
#[test]
fn index_value_cross_domain_comparison_and_ranges_are_typed_errors() {
assert_eq!(
IndexValue::Int64(0).compare(&IndexValue::UInt64(0)),
Err(IndexValueError::DomainMismatch {
left: "int64",
right: "uint64"
})
);
let kind = IndexKind::UInt64 {
index_granularity: NonZeroU64::new(1).unwrap(),
};
assert!(matches!(
validate_index_range(&kind, &IndexValue::Int64(0), &IndexValue::Int64(1)),
Err(IndexValueError::KindMismatch { .. })
));
assert!(matches!(
validate_index_range(&kind, &IndexValue::UInt64(1), &IndexValue::UInt64(1)),
Err(IndexValueError::InvalidRange { .. })
));
}
#[test]
fn time_index_granularity_parse_accepts_basic_units() {
let cases = [
("1s", TimeIndexGranularity::Seconds(1)),
("2m", TimeIndexGranularity::Minutes(2)),
("3h", TimeIndexGranularity::Hours(3)),
("4d", TimeIndexGranularity::Days(4)),
];
for (input, expected) in cases {
assert_eq!(input.parse::<TimeIndexGranularity>().unwrap(), expected);
}
}
#[test]
fn time_index_granularity_parse_accepts_aliases_case_and_whitespace() {
let cases = [
("1sec", TimeIndexGranularity::Seconds(1)),
("1secs", TimeIndexGranularity::Seconds(1)),
("1second", TimeIndexGranularity::Seconds(1)),
("1seconds", TimeIndexGranularity::Seconds(1)),
("1min", TimeIndexGranularity::Minutes(1)),
("1mins", TimeIndexGranularity::Minutes(1)),
("1minute", TimeIndexGranularity::Minutes(1)),
("1minutes", TimeIndexGranularity::Minutes(1)),
("1hr", TimeIndexGranularity::Hours(1)),
("1hrs", TimeIndexGranularity::Hours(1)),
("1hour", TimeIndexGranularity::Hours(1)),
("1hours", TimeIndexGranularity::Hours(1)),
("1day", TimeIndexGranularity::Days(1)),
("1days", TimeIndexGranularity::Days(1)),
("1H", TimeIndexGranularity::Hours(1)),
("1MiN", TimeIndexGranularity::Minutes(1)),
(" 2h", TimeIndexGranularity::Hours(2)),
("3d ", TimeIndexGranularity::Days(3)),
(" 4m ", TimeIndexGranularity::Minutes(4)),
("1 h", TimeIndexGranularity::Hours(1)),
];
for (input, expected) in cases {
assert_eq!(input.parse::<TimeIndexGranularity>().unwrap(), expected);
}
}
#[test]
fn time_index_granularity_parse_rejects_empty_or_whitespace() {
let cases = ["", " ", "\n\t"];
for input in cases {
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(matches!(err, ParseTimeIndexGranularityError::Empty));
}
}
#[test]
fn time_index_granularity_parse_rejects_missing_number() {
let cases = ["h", " hr", "day", "abcmin"];
for input in cases {
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(
matches!(err, ParseTimeIndexGranularityError::MissingNumber { .. }),
"expected MissingNumber for {input:?}, got {err:?}"
);
}
}
#[test]
fn time_index_granularity_parse_rejects_missing_unit() {
let cases = ["1", " 42 "];
for input in cases {
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(
matches!(err, ParseTimeIndexGranularityError::MissingUnit { .. }),
"expected MissingUnit for {input:?}, got {err:?}"
);
}
}
#[test]
fn time_index_granularity_parse_rejects_invalid_number() {
let cases = ["1.5h", "1_000s"];
for input in cases {
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(
matches!(err, ParseTimeIndexGranularityError::InvalidNumber { .. }),
"expected InvalidNumber for {input:?}, got {err:?}"
);
}
}
#[test]
fn time_index_granularity_parse_rejects_non_positive() {
let cases = ["0s", "0m"];
for input in cases {
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(
matches!(
err,
ParseTimeIndexGranularityError::NonPositive { value: 0, .. }
),
"expected NonPositive for {input:?}, got {err:?}"
);
}
}
#[test]
fn time_index_granularity_parse_rejects_too_large() {
let too_large = (u32::MAX as u64 + 1).to_string();
let input = format!("{too_large}h");
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(
matches!(err, ParseTimeIndexGranularityError::TooLarge { value, .. } if value == u32::MAX as u64 + 1),
"expected TooLarge for {input:?}, got {err:?}"
);
}
#[test]
fn time_index_granularity_parse_rejects_unknown_units() {
let cases = ["1w", "1ms", "1mo", "10msec"];
for input in cases {
let err = input.parse::<TimeIndexGranularity>().unwrap_err();
assert!(
matches!(err, ParseTimeIndexGranularityError::UnknownUnit { .. }),
"expected UnknownUnit for {input:?}, got {err:?}"
);
}
}
#[test]
fn time_index_granularity_parse_matches_from_str() {
let via_method = TimeIndexGranularity::parse("5m").unwrap();
let via_trait: TimeIndexGranularity = "5m".parse().unwrap();
assert_eq!(via_method, via_trait);
}
}