Expose read_segment

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wal-rus"
-version = "0.3.4"
+version = "0.3.5"
 edition = "2024"
 rust-version = "1.90"
 description = "Rust port of wal-g for PostgreSQL, optimized for no-overcommit hosts"

src/pg/wal/fetch.rs

@@ -209,31 +209,73 @@ pub(super) async fn download_to_running(
     Ok(true)
 }
 
-async fn find_object(
-    storage: &dyn crate::storage::Storage,
+/// Keys to try for `name`, configured compression first
+fn candidate_keys(
     name: &str,
     preferred: compression::Method,
-) -> Result<Option<(String, compression::Method)>> {
+) -> impl Iterator<Item = (String, compression::Method)> {
     let preferred_ext = preferred.extension();
-    let mut order: Vec<&str> = vec![preferred_ext];
-    for e in CANDIDATE_EXTS {
-        if !order.contains(e) {
-            order.push(e);
-        }
-    }
+    std::iter::once(preferred_ext)
+        .chain(
+            CANDIDATE_EXTS
+                .iter()
+                .copied()
+                .filter(move |e| *e != preferred_ext),
+        )
+        .map(move |ext| {
+            let key = if ext.is_empty() {
+                format!("{}/{}", pg::WAL_FOLDER, name)
+            } else {
+                format!("{}/{}.{}", pg::WAL_FOLDER, name, ext)
+            };
+            let method =
+                compression::Method::from_extension(ext).unwrap_or(compression::Method::None);
+            (key, method)
+        })
+}
 
-    for ext in order {
-        let key = if ext.is_empty() {
-            format!("{}/{}", pg::WAL_FOLDER, name)
-        } else {
-            format!("{}/{}.{}", pg::WAL_FOLDER, name, ext)
+/// Read one archived WAL object whole, over the same throttle, decrypt and
+/// decompress chain [`handle`] uses, without staging it on disk.
+///
+/// Unlike [`handle`], fetches each candidate directly instead of probing for
+/// it: a bucket written under one compression costs one request, not an
+/// existence check per extension
+pub async fn read_segment(
+    settings: &Settings,
+    storage: &DynStorage,
+    name: &str,
+) -> Result<Vec<u8>> {
+    let preferred = if is_history_filename(name) {
+        compression::Method::None
+    } else {
+        settings.compression
+    };
+    for (key, method) in candidate_keys(name, preferred) {
+        let body = match storage.get(&key).await {
+            Ok(body) => body,
+            Err(StorageError::NotFound(_)) => continue,
+            Err(e) => return Err(anyhow::Error::new(e).context(format!("get {key}"))),
         };
+        let mut decoded =
+            compression::decode(method, settings.decrypt(settings.throttle_network(body)));
+        let mut bytes = Vec::new();
+        decoded
+            .read_to_end(&mut bytes)
+            .await
+            .with_context(|| format!("read {key}"))?;
+        return Ok(bytes);
+    }
+    Err(ArchiveNotFound(name.to_string()).into())
+}
+
+async fn find_object(
+    storage: &dyn crate::storage::Storage,
+    name: &str,
+    preferred: compression::Method,
+) -> Result<Option<(String, compression::Method)>> {
+    for (key, method) in candidate_keys(name, preferred) {
         match storage.exists(&key).await {
-            Ok(true) => {
-                let m =
-                    compression::Method::from_extension(ext).unwrap_or(compression::Method::None);
-                return Ok(Some((key, m)));
-            }
+            Ok(true) => return Ok(Some((key, method))),
             Ok(false) => continue,
             Err(StorageError::NotFound(_)) => continue,
             Err(e) => return Err(e.into()),

tests/wal_roundtrip.rs

@@ -793,3 +793,37 @@ async fn ciphertext_overhead_matches_libsodium_layout() {
     let expected = 24 + (8192 + 17) + (2048 + 17);
     assert_eq!(stored_len, expected, "wire layout drift");
 }
+
+#[tokio::test]
+async fn read_segment_returns_bytes_across_compressions() {
+    let dir = tempfile::tempdir().unwrap();
+    let storage_dir = dir.path().join("storage");
+    let stage = dir.path().join("stage");
+    std::fs::create_dir_all(&stage).unwrap();
+    let name = "000000010000000000000009";
+    let src = stage.join(name);
+    std::fs::write(&src, b"in-memory wal").unwrap();
+
+    let store = Arc::new(FsStorage::new(&storage_dir).unwrap());
+    let pushed = settings_for(storage_dir.to_str().unwrap(), Method::None);
+    wal::push::handle(&pushed, store.clone(), &src)
+        .await
+        .unwrap();
+
+    // Bucket written under another compression still reads back
+    let reading = settings_for(storage_dir.to_str().unwrap(), Method::Zstd);
+    let bytes = wal::fetch::read_segment(&reading, &(store.clone() as _), name)
+        .await
+        .unwrap();
+    assert_eq!(bytes, b"in-memory wal");
+
+    let missing = wal::fetch::read_segment(&reading, &(store as _), "000000010000000000000010")
+        .await
+        .unwrap_err();
+    assert!(
+        missing
+            .downcast_ref::<wal::fetch::ArchiveNotFound>()
+            .is_some(),
+        "{missing:#}"
+    );
+}