perf(sessions): batch canonical observation body reads

crates/tracedecay-global-db/src/observation/canonical_bodies.rs

@@ -1,4 +1,5 @@
 use tracedecay_runtime_core::db::engine::{Executor, QueryExecutor, params};
+use tracedecay_store::canonical_body::LOAD_CANONICAL_BODIES_SQL;
 use tracedecay_store::{
     CANONICAL_BODIES_TABLE_SQL, CanonicalBodyError, INLINE_BODY_BYTES, LOAD_CANONICAL_BODY_SQL,
     StoredCanonicalBody, UPSERT_CANONICAL_BODY_SQL, collect_body_refs, parse_stored_observation,
@@ -112,8 +113,36 @@ pub async fn decode_observation_json(
 ) -> tracedecay_store::ObservationStoreResult<tracedecay_domain::DurableObservationV1> {
     let hashes = collect_body_refs(observation_json).map_err(|error| storage(operation, error))?;
     let mut bodies = std::collections::HashMap::new();
-    for hash in hashes {
-        bodies.insert(hash.clone(), load_canonical_body(conn, &hash).await?);
+    if !hashes.is_empty() {
+        const READ_OPERATION: &str = "read canonical observation body";
+        let encoded = serde_json::to_string(&hashes).map_err(|error| storage(operation, error))?;
+        let mut rows = conn
+            .query(LOAD_CANONICAL_BODIES_SQL, params![encoded])
+            .await
+            .map_err(|error| storage(READ_OPERATION, error))?;
+        while let Some(row) = rows
+            .next()
+            .await
+            .map_err(|error| storage(READ_OPERATION, error))?
+        {
+            let hash: String = row.get(0).map_err(|error| storage(READ_OPERATION, error))?;
+            let encoding = row
+                .get::<Option<String>>(1)
+                .map_err(|error| storage(READ_OPERATION, error))?
+                .ok_or_else(|| {
+                    storage(
+                        READ_OPERATION,
+                        CanonicalBodyError::Missing {
+                            content_hash: hash.clone(),
+                        },
+                    )
+                })?;
+            let blob: Vec<u8> = row.get(2).map_err(|error| storage(READ_OPERATION, error))?;
+            let uncompressed = row.get(3).map_err(|error| storage(READ_OPERATION, error))?;
+            let bytes = unpack_body(&hash, &encoding, &blob, uncompressed)
+                .map_err(|error| storage(READ_OPERATION, error))?;
+            bodies.insert(hash, bytes);
+        }
     }
     parse_stored_observation(observation_json, |hash| {
         bodies
@@ -197,6 +226,18 @@ async fn table_exists(
 
 #[cfg(test)]
 mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    use tracedecay_domain::{
+        ComponentVersion, DurableObservationV1, ObservationId, ObservationIdentityMaterialV1,
+        ObservationOrderingDomainV1, ObservationScopeV1, ObservationSourceGenerationV1,
+        ObservationSourceIdentityV1, ObservationSourceRangeV1, PayloadReferenceV1,
+        RetentionClass, SanitizationReceiptId, SanitizationReceiptRefV1,
+        SanitizationReceiptV1, SanitizerDispositionV1, SensitivityV1, SessionId,
+    };
+    use tracedecay_runtime_core::db::engine::{IntoParams, ReadConnection, Rows};
+    use tracedecay_store::slim_stored_json;
+
     use super::*;
     use tempfile::TempDir;
     use tracedecay_runtime_core::db::engine::TestConnection;
@@ -284,4 +325,140 @@ mod tests {
             "CAS rows keep neither the body nor a full placeholder"
         );
     }
+
+    fn body_observation(payload: serde_json::Value) -> DurableObservationV1 {
+        let identity = ObservationIdentityMaterialV1::for_native_record(
+            ObservationSourceIdentityV1::new(SessionId::new("session.bodies").unwrap())
+                .unwrap(),
+            ObservationScopeV1::Profile,
+            ObservationSourceGenerationV1::new(1).unwrap(),
+            ObservationSourceRangeV1::new(12, 34).unwrap(),
+            ObservationOrderingDomainV1::FileBytes,
+            ObservationId::new("record.bodies").unwrap(),
+        )
+        .unwrap();
+        let receipt = SanitizationReceiptV1::new(
+            SanitizationReceiptRefV1::new(
+                SanitizationReceiptId::new("receipt.bodies").unwrap(),
+                ComponentVersion::new("sanitizer.fixture.v1").unwrap(),
+            )
+            .unwrap(),
+            SanitizerDispositionV1::Accepted,
+            SensitivityV1::NonSensitive,
+            Some(PayloadReferenceV1::for_payload(&payload).unwrap()),
+        )
+        .unwrap();
+        DurableObservationV1::new(
+            identity,
+            receipt,
+            RetentionClass::new("transcript.fixture").unwrap(),
+            payload,
+        )
+        .unwrap()
+    }
+
+    struct CountedRead {
+        connection: ReadConnection,
+        queries: AtomicUsize,
+    }
+
+    impl QueryExecutor for CountedRead {
+        async fn query<P>(
+            &self,
+            sql: &str,
+            params: P,
+        ) -> tracedecay_runtime_core::db::engine::Result<Rows>
+        where
+            P: IntoParams,
+        {
+            self.queries.fetch_add(1, Ordering::Relaxed);
+            self.connection.query(sql, params).await
+        }
+    }
+
+    #[asupersync::test]
+    async fn observation_body_hydration_batches_reads_and_preserves_verification() {
+        let tmp = TempDir::new().unwrap();
+        let conn = TestConnection::open(&tmp.path().join("observation-bodies.db"));
+        ensure_canonical_bodies_table(&conn).await.unwrap();
+        let observation = body_observation(serde_json::json!({
+            "first": "alpha ".repeat(INLINE_BODY_BYTES),
+            "second": "beta ".repeat(INLINE_BODY_BYTES),
+            "same_as_first": "alpha ".repeat(INLINE_BODY_BYTES),
+        }));
+        let expected = serde_json::to_value(&observation).unwrap();
+        let (encoded, bodies) = slim_stored_json(&expected.to_string()).unwrap();
+        persist_bodies(&conn, &bodies).await.unwrap();
+        let hashes = collect_body_refs(&encoded).unwrap();
+        assert_eq!(hashes.len(), 2, "repeated payloads share one body");
+        // An unrelated corrupt row must not be read by the bounded lookup.
+        conn.execute(
+            UPSERT_CANONICAL_BODY_SQL,
+            params!["unrelated", "identity", b"wrong".as_slice(), 5_i64],
+        )
+        .await
+        .unwrap();
+        let counted = CountedRead {
+            connection: conn.read_only(),
+            queries: AtomicUsize::new(0),
+        };
+        let decoded = decode_observation_json(&counted, &encoded, "hydrate fixture")
+            .await
+            .unwrap();
+        assert_eq!(serde_json::to_value(decoded).unwrap(), expected);
+        assert_eq!(counted.queries.load(Ordering::Relaxed), 1);
+
+        let inline = body_observation(serde_json::json!({"text": "small"}));
+        let encoded_inline = serde_json::to_string(&inline).unwrap();
+        let decoded = decode_observation_json(&counted, &encoded_inline, "hydrate fixture")
+            .await
+            .unwrap();
+        assert_eq!(
+            serde_json::to_value(decoded).unwrap(),
+            serde_json::to_value(inline).unwrap()
+        );
+        assert_eq!(
+            counted.queries.load(Ordering::Relaxed),
+            1,
+            "inline rows issue no body query"
+        );
+
+        conn.execute(
+            "DELETE FROM session_canonical_bodies WHERE content_hash = ?1",
+            params![hashes[0].as_str()],
+        )
+        .await
+        .unwrap();
+        let error = decode_observation_json(&counted, &encoded, "hydrate fixture")
+            .await
+            .unwrap_err();
+        let tracedecay_store::ObservationStoreError::Storage { source, .. } = error else {
+            panic!("expected missing body refusal");
+        };
+        assert!(matches!(
+            source.downcast_ref::<CanonicalBodyError>(),
+            Some(CanonicalBodyError::Missing { content_hash }) if content_hash == &hashes[0]
+        ));
+        assert_eq!(counted.queries.load(Ordering::Relaxed), 2);
+
+        persist_bodies(&conn, &bodies).await.unwrap();
+        conn.execute(
+            "UPDATE session_canonical_bodies SET uncompressed_bytes = uncompressed_bytes + 1
+             WHERE content_hash = ?1",
+            params![hashes[1].as_str()],
+        )
+        .await
+        .unwrap();
+        let error = decode_observation_json(&counted, &encoded, "hydrate fixture")
+            .await
+            .unwrap_err();
+        let tracedecay_store::ObservationStoreError::Storage { source, .. } = error else {
+            panic!("expected corrupt body refusal");
+        };
+        assert!(matches!(
+            source.downcast_ref::<CanonicalBodyError>(),
+            Some(CanonicalBodyError::Corrupt { content_hash, .. }) if content_hash == &hashes[1]
+        ));
+        assert_eq!(counted.queries.load(Ordering::Relaxed), 3);
+    }
 }