timelake/
upgrade.rs

1//! Bringing a table to the current format (`timelake upgrade`) and undoing it (`rollback`).
2//!
3//! Every step is an ordinary commit, so nothing is lost: the previous snapshots stay readable,
4//! and the upgrade first sets `timelake.keep-snapshots-until` so none of them expires for a while
5//! (7 days by default). Within that time, [`Table::rollback`] returns the table to the snapshot it
6//! had before the upgrade.
7
8use std::collections::HashMap;
9use std::fmt;
10use std::time::Duration;
11
12use arrow::compute::concat_batches;
13use iceberg::spec::{MAIN_BRANCH, SnapshotReference, SnapshotRetention, TableMetadataBuilder};
14use timelake_format::head::NO_HEAD;
15use timelake_format::{FORMAT_VERSION, FormatVersion};
16
17use crate::commit::{Change, commit, publish};
18use crate::error::{Error, Result};
19use crate::lake::metadata_uri;
20use crate::maintain::{MaintainOptions, compaction_plan};
21use crate::schema::align_batch;
22use crate::storage::Uploads;
23use crate::table::Table;
24use crate::{Session, head, property, summary};
25
26/// Options for [`Table::upgrade`].
27#[derive(Clone, Debug)]
28pub struct UpgradeOptions {
29    /// Only work out what would change.
30    pub dry_run: bool,
31    /// Rewrite files that are unsorted or small, as `maintain` does.
32    pub compact: bool,
33    /// How long no snapshot may expire, so the upgrade can be rolled back.
34    pub keep_snapshots: Duration,
35}
36
37impl Default for UpgradeOptions {
38    fn default() -> Self {
39        Self { dry_run: false, compact: true, keep_snapshots: Duration::from_secs(7 * 24 * 3600) }
40    }
41}
42
43/// One change an upgrade makes, in the order it makes them.
44#[derive(Clone, Debug, PartialEq, Eq)]
45#[non_exhaustive]
46pub enum UpgradeStep {
47    /// Record the table format version, and keep every snapshot until `keep_until_ms`.
48    SetFormatVersion {
49        /// The format before.
50        from: FormatVersion,
51        /// The format after: [`FORMAT_VERSION`].
52        to: FormatVersion,
53    },
54    /// Build the last-value index the table lacks (a time-series table written by another engine).
55    BuildLastValueIndex,
56    /// Rewrite unsorted or small files.
57    Compact {
58        /// Partitions rewritten.
59        partitions: usize,
60        /// Data files replaced.
61        files: usize,
62    },
63    /// Write the head object of the final version.
64    WriteHead,
65}
66
67impl fmt::Display for UpgradeStep {
68    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69        match self {
70            Self::SetFormatVersion { from, to } => write!(f, "set table format {from} → {to}"),
71            Self::BuildLastValueIndex => f.write_str("build the last-value index"),
72            Self::Compact { partitions, files } => {
73                write!(f, "rewrite {files} unsorted or small files in {partitions} partitions")
74            }
75            Self::WriteHead => f.write_str("write the head object"),
76        }
77    }
78}
79
80/// What [`Table::upgrade`] did (or, with `dry_run`, would do).
81#[derive(Clone, Debug)]
82pub struct UpgradeReport {
83    /// The steps, in order; empty if the table was already current.
84    pub steps: Vec<UpgradeStep>,
85    /// The version before the upgrade.
86    pub version_before: u64,
87    /// The snapshot before the upgrade: pass it to [`Table::rollback`] to undo it.
88    pub snapshot_before: Option<i64>,
89    /// No snapshot expires before this time (milliseconds since the epoch), unless the upgrade
90    /// had nothing to do.
91    pub keep_snapshots_until_ms: Option<i64>,
92}
93
94impl Table {
95    /// Brings the table to the current format ([`FORMAT_VERSION`]) and layout, without losing
96    /// data: every step is a commit, and no snapshot expires for `opts.keep_snapshots`, so
97    /// [`Table::rollback`] to `report.snapshot_before` undoes it. Returns the new version.
98    pub async fn upgrade(&self, opts: &UpgradeOptions) -> Result<(Table, UpgradeReport)> {
99        let table = self.refresh().await?;
100        let steps = plan(&table, opts).await?;
101        let mut report = UpgradeReport {
102            steps,
103            version_before: table.version(),
104            snapshot_before: table.metadata().current_snapshot_id(),
105            keep_snapshots_until_ms: None,
106        };
107        if opts.dry_run || report.steps.is_empty() {
108            return Ok((table, report));
109        }
110        if report.steps == [UpgradeStep::WriteHead] {
111            head::write(&table).await?;
112            return Ok((table, report));
113        }
114        let keep_ms = i64::try_from(opts.keep_snapshots.as_millis()).unwrap_or(i64::MAX);
115        let until = chrono::Utc::now().timestamp_millis().saturating_add(keep_ms);
116        report.keep_snapshots_until_ms = Some(until);
117
118        // Protect the snapshots before anything else commits, and record the format.
119        let mut table = set_properties(
120            &table,
121            &[
122                (property::FORMAT_VERSION, FORMAT_VERSION.to_string()),
123                (property::KEEP_SNAPSHOTS_UNTIL, until.to_string()),
124            ],
125        )
126        .await?;
127        for step in &report.steps {
128            table = match step {
129                UpgradeStep::BuildLastValueIndex => build_index(&table).await?,
130                UpgradeStep::Compact { .. } => {
131                    let o = MaintainOptions { retention: false, expire: false, ..Default::default() };
132                    table.maintain(o).await?.0
133                }
134                UpgradeStep::SetFormatVersion { .. } | UpgradeStep::WriteHead => table,
135            };
136        }
137        if report.steps.contains(&UpgradeStep::WriteHead) {
138            head::write(&table).await?;
139        }
140        Ok((table, report))
141    }
142
143    /// Makes `snapshot_id`, an earlier snapshot of this table, current again: a new version whose
144    /// data is that snapshot's. Later snapshots stay in the history until they expire, so a
145    /// rollback can be rolled back too. Returns the new version.
146    pub async fn rollback(&self, snapshot_id: i64) -> Result<Table> {
147        let mut base = self.refresh().await?;
148        loop {
149            if base.metadata().snapshot_by_id(snapshot_id).is_none() {
150                return Err(Error::NotFound(format!("{}: snapshot {snapshot_id}", base.name())));
151            }
152            let previous = Some(metadata_uri(base.location(), base.version()));
153            let reference = SnapshotReference::new(snapshot_id, SnapshotRetention::branch(None, None, None));
154            let metadata = TableMetadataBuilder::new_from_metadata(base.metadata().clone(), previous)
155                .set_ref(MAIN_BRANCH, reference)?
156                .build()?
157                .metadata;
158            match publish(&base, metadata, None).await? {
159                Some(table) => return Ok(table),
160                None => base = base.refresh().await?,
161            }
162        }
163    }
164}
165
166/// The steps `table` needs.
167async fn plan(table: &Table, opts: &UpgradeOptions) -> Result<Vec<UpgradeStep>> {
168    let c = table.contract();
169    let mut steps = vec![];
170    if c.format_version < FORMAT_VERSION {
171        steps.push(UpgradeStep::SetFormatVersion { from: c.format_version, to: FORMAT_VERSION });
172    }
173    let has_rows = table.summary_value(summary::TOTAL_RECORDS).is_some_and(|n| n != "0");
174    if c.is_time_series() && has_rows && table.last_value_index().is_none() {
175        steps.push(UpgradeStep::BuildLastValueIndex);
176    }
177    if opts.compact {
178        let o = MaintainOptions { retention: false, expire: false, ..Default::default() };
179        let compaction = compaction_plan(table, &o).await?;
180        if !compaction.is_empty() {
181            let files = compaction.iter().map(|(_, files)| files.len()).sum();
182            steps.push(UpgradeStep::Compact { partitions: compaction.len(), files });
183        }
184    }
185    let head_current = head::stored_version(table).await? == Some(table.version());
186    if c.head_format != NO_HEAD && (!steps.is_empty() || !head_current) {
187        steps.push(UpgradeStep::WriteHead);
188    }
189    Ok(steps)
190}
191
192/// Commits new table properties (no new snapshot).
193pub(crate) async fn set_properties(table: &Table, properties: &[(&str, String)]) -> Result<Table> {
194    let mut base = table.clone();
195    loop {
196        let previous = Some(metadata_uri(base.location(), base.version()));
197        let props = properties.iter().map(|(k, v)| ((*k).to_string(), v.clone())).collect();
198        let metadata = TableMetadataBuilder::new_from_metadata(base.metadata().clone(), previous)
199            .set_properties(props)?
200            .build()?
201            .metadata;
202        match publish(&base, metadata, base.known_files().cloned()).await? {
203            Some(table) => return Ok(table),
204            None => base = base.refresh().await?,
205        }
206    }
207}
208
209/// Builds the newest row per series from the data and commits it as the last-value index.
210async fn build_index(table: &Table) -> Result<Table> {
211    let session = Session::new(table.lake().clone());
212    let batches = session.latest(table.name().clone()).await?.collect().await?;
213    let schema = table.schema();
214    let aligned = batches.iter().map(|b| align_batch(b, &schema)).collect::<Result<Vec<_>>>()?;
215    let latest = concat_batches(&schema, &aligned)?;
216    Ok(commit(table, Change::Index { latest }, HashMap::new(), Uploads::default()).await?.0)
217}
218
219#[cfg(test)]
220mod tests {
221    use std::collections::HashMap;
222
223    use super::*;
224    use crate::testing::{T0, append_without_index, batch, faulty_lake, opts};
225    use crate::{AppendOptions, MaintainOptions};
226
227    fn rows(t: &Table) -> u64 {
228        t.summary_value(summary::TOTAL_RECORDS).and_then(|v| v.parse().ok()).unwrap_or(0)
229    }
230
231    /// A table as an older Timelake or another engine leaves it: format 0.1.0, no last-value index,
232    /// no head, small files.
233    async fn old_table() -> (crate::Lake, Table) {
234        let (lake, _store) = faulty_lake();
235        let mut o = opts();
236        o.properties = HashMap::from([
237            (property::FORMAT_VERSION.into(), FormatVersion::UNVERSIONED.to_string()),
238            (property::HEAD_FORMAT.into(), NO_HEAD.into()),
239        ]);
240        let t = lake.create_table("iot.r", &crate::testing::schema(), o).await.unwrap();
241        let t = append_without_index(&t, &[batch(&[(T0, 1, 1.0), (T0 + 1, 2, 2.0)])]).await;
242        let t = append_without_index(&t, &[batch(&[(T0 + 2, 1, 3.0)])]).await;
243        // Writers from now on write heads, as a current Timelake does by default.
244        let props = [(property::HEAD_FORMAT, "arrow".to_string())];
245        let t = set_properties(&t, &props).await.unwrap();
246        crate::flush().await;
247        lake.storage().delete(&head::uri(t.location())).await.unwrap();
248        (lake, t)
249    }
250
251    #[tokio::test]
252    async fn upgrade_brings_a_table_to_the_current_format_and_rolls_back() {
253        let (_lake, t) = old_table().await;
254        let dry = UpgradeOptions { dry_run: true, ..Default::default() };
255        let (same, plan) = t.upgrade(&dry).await.unwrap();
256        assert_eq!(
257            plan.steps,
258            [
259                UpgradeStep::SetFormatVersion { from: FormatVersion::UNVERSIONED, to: FORMAT_VERSION },
260                UpgradeStep::BuildLastValueIndex,
261                UpgradeStep::Compact { partitions: 1, files: 2 },
262                UpgradeStep::WriteHead,
263            ]
264        );
265        assert_eq!(same.version(), t.version(), "a dry run commits nothing");
266
267        let (up, report) = t.upgrade(&UpgradeOptions::default()).await.unwrap();
268        assert_eq!(report.steps, plan.steps);
269        assert_eq!(up.contract().format_version, FORMAT_VERSION);
270        assert!(up.last_value_index().is_some());
271        assert_eq!((up.files().await.unwrap().len(), rows(&up)), (1, 3));
272        assert_eq!(head::stored_version(&up).await.unwrap(), Some(up.version()));
273        assert!(up.contract().keep_snapshots_until.is_some());
274
275        // Protected snapshots don't expire, so the rollback finds the old files.
276        let (up, m) = up.maintain(MaintainOptions::default()).await.unwrap();
277        assert_eq!(m.snapshots_expired, 0);
278        let back = up.rollback(report.snapshot_before.unwrap()).await.unwrap();
279        assert_eq!((back.files().await.unwrap().len(), rows(&back)), (2, 3));
280        assert!(back.last_value_index().is_none());
281
282        // Upgrading again redoes the work; a current table needs nothing.
283        let (again, _) = back.upgrade(&UpgradeOptions::default()).await.unwrap();
284        let (_, nothing) = again.upgrade(&UpgradeOptions::default()).await.unwrap();
285        assert!(nothing.steps.is_empty(), "{:?}", nothing.steps);
286    }
287
288    /// After a rollback, a batch from the abandoned snapshots is no longer a duplicate.
289    #[tokio::test]
290    async fn rolled_back_batches_can_be_appended_again() {
291        let (lake, _store) = faulty_lake();
292        let t = lake.create_table("iot.r", &crate::testing::schema(), opts()).await.unwrap();
293        let o = |id: &str| AppendOptions { batch_id: Some(id.into()), maintain: Some(false), ..Default::default() };
294        let (t, _) = t.append(&[batch(&[(T0, 1, 1.0)])], o("a")).await.unwrap();
295        let before = t.metadata().current_snapshot_id().unwrap();
296        let (t, _) = t.append(&[batch(&[(T0 + 1, 1, 2.0)])], o("b")).await.unwrap();
297        let t = t.rollback(before).await.unwrap();
298        let (t, r) = t.append(&[batch(&[(T0 + 1, 1, 2.0)])], o("b")).await.unwrap();
299        assert_eq!((r.rows, rows(&t)), (1, 2));
300        let (_, r) = t.append(&[batch(&[(T0, 1, 1.0)])], o("a")).await.unwrap();
301        assert!(r.commit.duplicate);
302    }
303}
timelake/upgrade.rs

timelake/
upgrade.rs
