fix: get next_sstable_id

src/db/builder.rs

@@ -20,8 +20,8 @@ use crate::{
     extractor::{KeyExtractError, KeyProjection, projection_for_columns},
     id::FileIdGenerator,
     manifest::{
-        ManifestError, ManifestFs, TableMeta, TonboManifest,
-        bootstrap::{TableSnapshot, ensure_manifest_dirs, init_fs_manifest},
+        ManifestError, ManifestFs, TableMeta, TonboManifest, VersionState,
+        bootstrap::{ensure_manifest_dirs, init_fs_manifest},
     },
     mode::{DynModeConfig, table_definition},
     ondisk::sstable::SsTableConfig,
@@ -1247,11 +1247,23 @@ where
         FS: ManifestFs<E>,
         <FS as fusio::fs::Fs>::File: fusio::durability::FileCommit,
     {
-        let snapshot = manifest
-            .snapshot_latest_with_fallback(table_meta.table_id, table_meta)
-            .await
-            .map_err(DbBuildError::Manifest)?;
-        let start_id = cfg.start_id.unwrap_or_else(|| next_sstable_id(&snapshot));
+        let start_id: u64 = {
+            let snapshot = manifest
+                .snapshot_latest_with_fallback(table_meta.table_id, table_meta)
+                .await
+                .map_err(DbBuildError::Manifest)?;
+
+            let versions = match snapshot.latest_version.as_ref() {
+                Some(v) if !v.ssts().is_empty() => std::slice::from_ref(v),
+                _ => &manifest
+                    .list_versions(table_meta.table_id, 0)
+                    .await
+                    .map_err(DbBuildError::Manifest)?,
+            };
+
+            cfg.start_id.unwrap_or_else(|| next_sstable_id(versions))
+        };
+
         Self::build_minor_compaction_state(layout, cfg, start_id, schema, extractor)
     }
 
@@ -1344,16 +1356,16 @@ where
     }
 }
 
-fn next_sstable_id(snapshot: &TableSnapshot) -> u64 {
-    snapshot
-        .latest_version
-        .as_ref()
-        .and_then(|version| {
+fn next_sstable_id(versions: &[VersionState]) -> u64 {
+    let max_id = versions
+        .iter()
+        .flat_map(|version| {
             version
                 .ssts()
                 .iter()
                 .flat_map(|level| level.iter().map(|entry| entry.sst_id().raw()))
-                .max()
         })
-        .map_or(1, |max_id| max_id.saturating_add(1))
+        .max();
+
+    max_id.map_or(1, |max_id| max_id.saturating_add(1))
 }

tests/next_sstable_id.rs

@@ -0,0 +1,79 @@
+use std::sync::Arc;
+
+use arrow_array::{Int64Array, RecordBatch, StringArray};
+use arrow_schema::{DataType, Field, Schema};
+use fusio::executor::NoopExecutor;
+use tonbo::db::DbBuilder;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_sstable_id_persistence_across_restarts() {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Utf8, false),
+        Field::new("value", DataType::Int64, false),
+    ]));
+
+    let db = DbBuilder::from_schema_key_name(schema.clone(), "id")
+        .unwrap()
+        .in_memory("test_db")
+        .unwrap()
+        .open_with_executor(Arc::new(NoopExecutor))
+        .await
+        .unwrap();
+
+    let batch1 = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(StringArray::from(vec!["key1", "key2"])),
+            Arc::new(Int64Array::from(vec![100, 200])),
+        ],
+    )
+    .unwrap();
+    db.ingest(batch1).await.unwrap();
+
+    let batch2 = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(StringArray::from(vec!["key3", "key4"])),
+            Arc::new(Int64Array::from(vec![300, 400])),
+        ],
+    )
+    .unwrap();
+    db.ingest(batch2).await.unwrap();
+
+    // Give db time to compact
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+    let batches = db.scan().collect().await.unwrap();
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(total_rows, 4, "Should have all 4 rows in the database");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_db_with_no_latest_version() {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Utf8, false),
+        Field::new("value", DataType::Int64, false),
+    ]));
+
+    let db = DbBuilder::from_schema_key_name(schema.clone(), "id")
+        .unwrap()
+        .in_memory("empty_db")
+        .unwrap()
+        .open_with_executor(Arc::new(NoopExecutor))
+        .await
+        .unwrap();
+
+    let batch = RecordBatch::try_new(
+        schema,
+        vec![
+            Arc::new(StringArray::from(vec!["first_key"])),
+            Arc::new(Int64Array::from(vec![1])),
+        ],
+    )
+    .unwrap();
+    db.ingest(batch).await.unwrap();
+
+    let batches = db.scan().collect().await.unwrap();
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(total_rows, 1);
+}