use crate::chunking::Bounds;
use crate::error::{CliError, CliResult};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum PartitionSpec {
Integer {
from: i64,
to: IntBound,
chunk_size: u64,
bounds: Bounds,
#[serde(default, skip_serializing_if = "Option::is_none")]
to_unbounded: Option<bool>,
},
Timestamp {
from: String,
to: String,
chunk_size: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
timezone: Option<String>,
},
Offset {
total: CountBound,
chunk_size: u64,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(untagged)]
pub enum IntBound {
Literal(i64),
Discovered(BoundProbe),
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(untagged)]
pub enum CountBound {
Literal(u64),
Discovered(BoundProbe),
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
#[serde(deny_unknown_fields)]
pub struct BoundProbe {
pub from_source: crate::config::ConnectorSpec,
pub value_path: String,
}
impl PartitionSpec {
pub fn validate(&self) -> CliResult<()> {
match self {
Self::Integer { chunk_size, .. } | Self::Offset { chunk_size, .. } => {
if *chunk_size == 0 {
return Err(CliError::Config(
"partition.chunk_size must be greater than 0 (unlike `batch_size`, \
0 is not a 'no chunking' sentinel here)"
.into(),
));
}
}
Self::Timestamp {
chunk_size,
timezone,
..
} => {
crate::chunking::parse_window(chunk_size)?;
if let Some(tz) = timezone {
tz.parse::<chrono_tz::Tz>().map_err(|_| {
CliError::Config(format!(
"'{tz}' is not a valid IANA timezone (e.g. UTC, America/New_York)"
))
})?;
}
}
}
Ok(())
}
pub fn token_names(&self) -> &'static [&'static str] {
match self {
Self::Integer { .. } => &["start", "end", "index", "id"],
Self::Timestamp { .. } => &[
"start",
"end",
"start_date",
"end_date",
"start_unix",
"end_unix",
"index",
"id",
],
Self::Offset { .. } => &["offset", "limit", "index", "id"],
}
}
pub fn kind_str(&self) -> &'static str {
match self {
Self::Integer { .. } => "integer",
Self::Timestamp { .. } => "timestamp",
Self::Offset { .. } => "offset",
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn parse(yaml: &str) -> Result<PartitionSpec, serde_yaml::Error> {
serde_yaml::from_str(yaml)
}
#[test]
fn integer_requires_bounds() {
let err = parse("kind: integer\nfrom: 0\nto: 100\nchunk_size: 10\n")
.expect_err("bounds must be required");
assert!(err.to_string().contains("bounds"), "{err}");
}
#[test]
fn integer_parses_with_bounds() {
let s =
parse("kind: integer\nfrom: 0\nto: 100\nchunk_size: 10\nbounds: inclusive\n").unwrap();
assert!(matches!(
s,
PartitionSpec::Integer {
from: 0,
to: IntBound::Literal(100),
chunk_size: 10,
bounds: Bounds::Inclusive,
to_unbounded: None
}
));
s.validate().unwrap();
}
#[test]
fn a_count_cannot_bound_an_id_range() {
let err = parse("kind: integer\nfrom: 0\ntotal: 1000\nchunk_size: 10\nbounds: inclusive\n")
.expect_err("total is not an integer-range field");
assert!(err.to_string().contains("total"), "{err}");
}
#[test]
fn an_id_bound_cannot_be_given_to_an_offset_range() {
let err = parse("kind: offset\nto: 1000\nchunk_size: 10\n")
.expect_err("to is not an offset field");
assert!(err.to_string().contains("to"), "{err}");
}
#[test]
fn offset_parses_and_needs_no_bounds() {
let s = parse("kind: offset\ntotal: 250\nchunk_size: 100\n").unwrap();
assert!(matches!(
s,
PartitionSpec::Offset {
total: CountBound::Literal(250),
chunk_size: 100
}
));
s.validate().unwrap();
}
#[test]
fn timestamp_parses_and_validates_window_and_timezone() {
let s = parse(
"kind: timestamp\nfrom: 2026-01-01\nto: 2026-02-01\nchunk_size: 1d\n\
timezone: America/New_York\n",
)
.unwrap();
s.validate().unwrap();
let bad_window = parse("kind: timestamp\nfrom: a\nto: b\nchunk_size: 1y\n").unwrap();
assert!(bad_window.validate().is_err(), "1y is not a valid window");
let bad_tz =
parse("kind: timestamp\nfrom: a\nto: b\nchunk_size: 1d\ntimezone: Mars/Olympus\n")
.unwrap();
assert!(
bad_tz.validate().is_err(),
"bogus timezone must be rejected"
);
}
#[test]
fn zero_chunk_size_is_rejected_with_the_batch_size_distinction_spelled_out() {
let s =
parse("kind: integer\nfrom: 0\nto: 10\nchunk_size: 0\nbounds: half_open\n").unwrap();
let err = s.validate().unwrap_err();
assert!(err.to_string().contains("greater than 0"), "{err}");
assert!(
err.to_string().contains("batch_size"),
"the message should distinguish it from the batch_size sentinel: {err}"
);
}
#[test]
fn unknown_fields_are_rejected() {
assert!(
parse("kind: offset\ntotal: 10\nchunk_size: 5\nchunck_size: 5\n").is_err(),
"a typo must not be silently ignored"
);
}
#[test]
fn token_names_match_the_kind() {
let int =
parse("kind: integer\nfrom: 0\nto: 1\nchunk_size: 1\nbounds: inclusive\n").unwrap();
assert!(int.token_names().contains(&"start"));
assert!(!int.token_names().contains(&"offset"));
let off = parse("kind: offset\ntotal: 1\nchunk_size: 1\n").unwrap();
assert!(off.token_names().contains(&"offset"));
assert!(!off.token_names().contains(&"start"));
}
}