use std::{collections::HashMap, path::Path};
#[cfg(feature = "test-counters")]
use std::cell::Cell;
#[cfg(feature = "test-counters")]
thread_local! {
static REBUILD_TABLE_STATE_COUNT: Cell<usize> = const { Cell::new(0) };
}
#[cfg(feature = "test-counters")]
pub fn rebuild_table_state_count() -> usize {
REBUILD_TABLE_STATE_COUNT.with(|c| c.get())
}
#[cfg(feature = "test-counters")]
pub fn reset_rebuild_table_state_count() {
REBUILD_TABLE_STATE_COUNT.with(|c| c.set(0));
}
use crate::{
metadata::{segments::cmp_segment_meta_by_time, table_metadata::TABLE_FORMAT_VERSION},
storage::normalize_relative_segment_path,
transaction_log::*,
};
fn validate_persisted_segment_path(path: &str) -> Result<(), CommitError> {
let (canonical, _) = match normalize_relative_segment_path(Path::new(path)) {
Ok(path) => path,
Err(source) => {
return CorruptStateSnafu {
msg: format!("Invalid persisted segment path {path:?}: {source}"),
}
.fail();
}
};
if canonical != path {
return CorruptStateSnafu {
msg: format!(
"Non-canonical persisted segment path {path:?}; canonical form is {canonical:?}"
),
}
.fail();
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableCoveragePointer {
pub bucket_spec: TimeBucket,
pub coverage_path: String,
pub version: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableState {
pub version: u64,
pub table_meta: TableMeta,
pub segments: HashMap<String, SegmentMeta>,
pub table_coverage: Option<TableCoveragePointer>,
}
impl TableState {
pub fn segments_sorted_by_time(&self) -> Vec<&SegmentMeta> {
let mut v: Vec<&SegmentMeta> = self.segments.values().collect();
v.sort_unstable_by(|a, b| cmp_segment_meta_by_time(a, b));
v
}
}
impl TransactionLogStore {
pub async fn rebuild_table_state(&self) -> Result<TableState, CommitError> {
#[cfg(feature = "test-counters")]
REBUILD_TABLE_STATE_COUNT.with(|c| c.set(c.get() + 1));
let current_version = self.load_current_version().await?;
if current_version == 0 {
return CorruptStateSnafu {
msg: "Cannot rebuild TableState: CURRENT is 0 (no commits)".to_string(),
}
.fail();
}
let mut table_meta: Option<TableMeta> = None;
let mut segments: HashMap<String, SegmentMeta> = HashMap::new();
let mut table_coverage: Option<TableCoveragePointer> = None;
for v in 1..=current_version {
let commit = self.load_commit(v).await?;
if commit.version != v {
return CorruptStateSnafu {
msg: format!(
"Commit version mismatch: expected {v}, found {} in payload",
commit.version
),
}
.fail();
}
for action in commit.actions {
match action {
LogAction::AddSegment(meta) => {
validate_persisted_segment_path(&meta.path)?;
if segments.contains_key(&meta.path) {
return CorruptStateSnafu {
msg: format!("Duplicate live segment path: {}", meta.path),
}
.fail();
}
segments.insert(meta.path.clone(), meta);
}
LogAction::RemoveSegment { path } => {
validate_persisted_segment_path(&path)?;
segments.remove(&path);
}
LogAction::UpdateTableMeta(delta) => {
if delta.format_version() != TABLE_FORMAT_VERSION {
return CorruptStateSnafu {
msg: format!(
"Unsupported table format version: expected {TABLE_FORMAT_VERSION}, found {}",
delta.format_version()
),
}
.fail();
}
table_meta = Some(delta);
}
LogAction::UpdateTableCoverage {
bucket_spec,
coverage_path,
} => {
table_coverage = Some(TableCoveragePointer {
bucket_spec,
coverage_path,
version: v,
})
}
}
}
}
let table_meta = table_meta.context(CorruptStateSnafu {
msg: format!("No TableMeta found in commits up to version {current_version}",),
})?;
Ok(TableState {
version: current_version,
table_meta,
segments,
table_coverage,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::layout;
use crate::storage::{StorageError, TableLocation};
use crate::transaction_log::{
FileFormat, LogAction, SegmentMeta, TableKind, TableMeta, TimeBucket, TimeIndexSpec,
TransactionLogStore,
};
use chrono::TimeZone;
use tempfile::TempDir;
type TestResult = Result<(), Box<dyn std::error::Error>>;
fn create_test_log_store() -> (TempDir, TransactionLogStore) {
let tmp = TempDir::new().expect("create temp dir");
let location = TableLocation::local(tmp.path());
let store = TransactionLogStore::new(location);
(tmp, store)
}
fn sample_table_meta() -> TableMeta {
TableMeta {
kind: TableKind::TimeSeries(TimeIndexSpec {
timestamp_column: "ts".to_string(),
entity_columns: vec!["symbol".to_string()],
bucket: TimeBucket::Minutes(1),
timezone: None,
}),
logical_schema: None,
created_at: chrono::Utc
.with_ymd_and_hms(2025, 1, 1, 0, 0, 0)
.single()
.expect("valid sample table metadata timestamp"),
format_version: TABLE_FORMAT_VERSION,
entity_identity: None,
}
}
fn sample_segment(id: &str) -> SegmentMeta {
SegmentMeta {
path: format!("data/{id}.parquet"),
format: FileFormat::Parquet,
ts_min: chrono::Utc
.with_ymd_and_hms(2025, 1, 1, 0, 0, 0)
.single()
.expect("valid sample segment ts_min"),
ts_max: chrono::Utc
.with_ymd_and_hms(2025, 1, 1, 1, 0, 0)
.single()
.expect("valid sample segment ts_max"),
row_count: 42,
file_size: None,
coverage_path: None,
}
}
fn segment_with_ts(id: &str, ts_min: i64, ts_max: i64) -> SegmentMeta {
SegmentMeta {
path: format!("data/{id}.parquet"),
format: FileFormat::Parquet,
ts_min: chrono::Utc.timestamp_opt(ts_min, 0).single().unwrap(),
ts_max: chrono::Utc.timestamp_opt(ts_max, 0).single().unwrap(),
row_count: 1,
file_size: None,
coverage_path: None,
}
}
#[test]
fn segments_sorted_by_time_orders_hashmap_deterministically() {
let mut segments = HashMap::new();
let seg_c = segment_with_ts("c", 10, 30);
let seg_a = segment_with_ts("a", 10, 20);
let seg_d = segment_with_ts("d", 5, 7);
let seg_b = segment_with_ts("b", 10, 20);
segments.insert(seg_c.path.clone(), seg_c);
segments.insert(seg_a.path.clone(), seg_a);
segments.insert(seg_d.path.clone(), seg_d);
segments.insert(seg_b.path.clone(), seg_b);
let state = TableState {
version: 3,
table_meta: sample_table_meta(),
segments,
table_coverage: None,
};
let ordered: Vec<(i64, i64, String)> = state
.segments_sorted_by_time()
.iter()
.map(|seg| {
(
seg.ts_min.timestamp(),
seg.ts_max.timestamp(),
seg.path.clone(),
)
})
.collect();
let mut expected = ordered.clone();
expected.sort();
assert_eq!(ordered, expected);
}
#[tokio::test]
async fn rebuild_table_state_happy_path() -> TestResult {
let (_tmp, store) = create_test_log_store();
let meta = sample_table_meta();
let seg1 = sample_segment("seg1");
let seg2 = sample_segment("seg2");
let v1 = store
.commit_with_expected_version(0, vec![LogAction::UpdateTableMeta(meta.clone())])
.await?;
let v2 = store
.commit_with_expected_version(
v1,
vec![
LogAction::AddSegment(seg1.clone()),
LogAction::AddSegment(seg2.clone()),
],
)
.await?;
let v3 = store
.commit_with_expected_version(
v2,
vec![LogAction::RemoveSegment {
path: seg1.path.clone(),
}],
)
.await?;
let state = store.rebuild_table_state().await?;
assert_eq!(state.version, v3);
assert_eq!(state.table_meta, meta);
assert!(state.segments.contains_key(&seg2.path));
assert!(!state.segments.contains_key(&seg1.path));
Ok(())
}
#[tokio::test]
async fn rebuild_table_state_errors_when_current_zero() {
let (_tmp, store) = create_test_log_store();
let err = store
.rebuild_table_state()
.await
.expect_err("expected error");
assert!(matches!(err, CommitError::CorruptState { .. }));
}
#[tokio::test]
async fn rebuild_table_state_errors_when_no_table_meta() -> TestResult {
let (_tmp, store) = create_test_log_store();
let seg = sample_segment("seg");
store
.commit_with_expected_version(0, vec![LogAction::AddSegment(seg.clone())])
.await?;
let err = store
.rebuild_table_state()
.await
.expect_err("expected error");
assert!(matches!(err, CommitError::CorruptState { .. }));
Ok(())
}
#[tokio::test]
async fn rebuild_table_state_rejects_old_format_version() -> TestResult {
let (_tmp, store) = create_test_log_store();
let mut meta = sample_table_meta();
meta.format_version = TABLE_FORMAT_VERSION - 1;
store
.commit_with_expected_version(0, vec![LogAction::UpdateTableMeta(meta)])
.await?;
let err = store
.rebuild_table_state()
.await
.expect_err("old format version should be rejected");
assert!(matches!(err, CommitError::CorruptState { .. }));
assert!(err.to_string().contains(&format!(
"expected {TABLE_FORMAT_VERSION}, found {}",
TABLE_FORMAT_VERSION - 1
)));
Ok(())
}
#[tokio::test]
async fn rebuild_table_state_rejects_noncanonical_segment_action_paths() -> TestResult {
for path in [
"",
"/data/seg.parquet",
"../data/seg.parquet",
"data/../seg.parquet",
r"data\seg.parquet",
"data//seg.parquet",
r"C:\data\seg.parquet",
"data/C:/seg.parquet",
"data/C:seg.parquet",
] {
let mut segment = sample_segment("seg");
segment.path = path.to_owned();
for action in [
LogAction::AddSegment(segment.clone()),
LogAction::RemoveSegment {
path: path.to_owned(),
},
] {
let (_tmp, store) = create_test_log_store();
store
.commit_with_expected_version(
0,
vec![LogAction::UpdateTableMeta(sample_table_meta()), action],
)
.await?;
let err = store
.rebuild_table_state()
.await
.expect_err("noncanonical segment action path should be rejected");
assert!(matches!(err, CommitError::CorruptState { .. }));
assert!(err.to_string().contains("segment path"), "{err}");
}
}
Ok(())
}
#[tokio::test]
async fn rebuild_table_state_fails_on_corrupt_commit_payload() -> TestResult {
let (tmp, store) = create_test_log_store();
let meta = sample_table_meta();
store
.commit_with_expected_version(0, vec![LogAction::UpdateTableMeta(meta)])
.await?;
let commit_path = tmp.path().join(layout::commit_rel_path(1));
tokio::fs::write(&commit_path, b"not-json").await?;
let err = store
.rebuild_table_state()
.await
.expect_err("expected error");
assert!(matches!(err, CommitError::CorruptState { .. }));
Ok(())
}
#[tokio::test]
async fn rebuild_table_state_fails_when_commit_missing() -> TestResult {
let (tmp, store) = create_test_log_store();
let meta = sample_table_meta();
store
.commit_with_expected_version(0, vec![LogAction::UpdateTableMeta(meta)])
.await?;
let commit_path = tmp.path().join(layout::commit_rel_path(1));
tokio::fs::remove_file(&commit_path).await?;
let err = store
.rebuild_table_state()
.await
.expect_err("expected error");
match err {
CommitError::Storage { source } => match source {
StorageError::NotFound { .. } => {}
other => panic!("unexpected storage error: {other:?}"),
},
other => panic!("expected storage error, got {other:?}"),
}
Ok(())
}
}