Support GetMemoriesById.

Change-Id: I5cd6cc37dc0e4fcf72306e380b7a971aeb0954f9
BUG: 482877275

Author: Changochen

oak_private_memory/app/handler.rs

@@ -179,6 +179,21 @@ impl SealedMemorySessionHandler {
         Ok(GetMemoryByIdResponse { memory, success })
     }
 
+    pub async fn get_memories_by_id_handler(
+        &self,
+        request: GetMemoriesByIdRequest,
+    ) -> tonic::Result<GetMemoriesByIdResponse> {
+        let mut mutex_guard = self.session_context().await;
+        let database =
+            &mut mutex_guard.as_mut().into_failed_precondition("call key sync first")?.database;
+
+        let (memories, not_found_ids) = database
+            .get_memories_by_id(request.ids, &request.result_mask)
+            .await
+            .into_internal_error("failed to get memories by id")?;
+        Ok(GetMemoriesByIdResponse { memories, not_found_ids })
+    }
+
     pub async fn reset_memory_handler(
         &self,
         _request: ResetMemoryRequest,
@@ -481,6 +496,9 @@ impl SealedMemorySessionHandler {
             sealed_memory_request::Request::DeleteMemoryRequest(request) => {
                 self.delete_memory_handler(request).await?.into_response()
             }
+            sealed_memory_request::Request::GetMemoriesByIdRequest(request) => {
+                self.get_memories_by_id_handler(request).await?.into_response()
+            }
         };
         let elapsed_time = start_time.elapsed().as_millis() as u64;
         self.metrics.record_latency(elapsed_time, metric_name);

oak_private_memory/app/packing.rs

@@ -125,3 +125,5 @@ impl_packing!(Response => GetMemoryByIdResponse);
 impl_packing!(Response => SearchMemoryResponse);
 impl_packing!(Response => DeleteMemoryResponse);
 impl_packing!(Response => UserRegistrationResponse);
+impl_packing!(Request => GetMemoriesByIdRequest);
+impl_packing!(Response => GetMemoriesByIdResponse);

oak_private_memory/database/database_with_cache.rs

@@ -166,6 +166,31 @@ impl DatabaseWithCache {
         }
     }
 
+    /// Returns memories for the given IDs.
+    /// Memories can be returned in an arbitrary order.
+    /// If any memory is not found, it will be skipped and its ID collected in
+    /// the second return value.
+    pub async fn get_memories_by_id(
+        &mut self,
+        ids: Vec<MemoryId>,
+        result_mask: &Option<ResultMask>,
+    ) -> anyhow::Result<(Vec<Memory>, Vec<MemoryId>)> {
+        let mut found_blob_ids: Vec<BlobId> = Vec::new();
+        let mut not_found_ids: Vec<MemoryId> = Vec::new();
+
+        for id in ids {
+            match self.meta_db().get_blob_id_by_memory_id(id.clone())? {
+                Some(blob_id) => found_blob_ids.push(blob_id),
+                None => not_found_ids.push(id),
+            }
+        }
+
+        let mut memories = self.cache.get_memories_by_blob_ids(&found_blob_ids).await?;
+        Self::apply_mask_to_memories(&mut memories, result_mask);
+
+        Ok((memories, not_found_ids))
+    }
+
     pub async fn reset_memory(&mut self) -> anyhow::Result<()> {
         let all_memory_ids = self.meta_db().get_all_memory_ids()?;
         if !all_memory_ids.is_empty() {

oak_private_memory/proto/build.rs

@@ -44,6 +44,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         "oak.private_memory.KeySyncResponse",
         "oak.private_memory.GetMemoryByIdRequest",
         "oak.private_memory.GetMemoryByIdResponse",
+        "oak.private_memory.GetMemoriesByIdRequest",
+        "oak.private_memory.GetMemoriesByIdResponse",
         "oak.private_memory.SearchMemoryRequest",
         "oak.private_memory.SearchMemoryResponse",
         "oak.private_memory.Embedding",

oak_private_memory/proto/prelude.rs

@@ -17,14 +17,14 @@ pub mod v1 {
     pub use crate::oak::private_memory::{
         AddMemoryRequest, AddMemoryResponse, DataBlob, DeleteMemoryRequest, DeleteMemoryResponse,
         Embedding, EmbeddingQuery, EmbeddingQueryMetricType, EncryptedDataBlob,
-        EncryptedMetadataBlob, EncryptedUserInfo, GetMemoriesRequest, GetMemoriesResponse,
-        GetMemoryByIdRequest, GetMemoryByIdResponse, InvalidRequestResponse, KeyDerivationInfo,
-        KeySyncRequest, KeySyncResponse, Memory, MemoryContent, MemoryField, MemoryValue,
-        PlainTextUserInfo, ResetMemoryRequest, ResetMemoryResponse, ResultMask, ScoreRange,
-        SealedMemoryCredentials, SealedMemoryRequest, SealedMemoryResponse,
-        SealedMemorySessionRequest, SealedMemorySessionResponse, SearchMemoryQuery,
-        SearchMemoryRequest, SearchMemoryResponse, SearchMemoryResultItem, UserDb,
-        UserRegistrationRequest, UserRegistrationResponse, WrappedDataEncryptionKey,
+        EncryptedMetadataBlob, EncryptedUserInfo, GetMemoriesByIdRequest, GetMemoriesByIdResponse,
+        GetMemoriesRequest, GetMemoriesResponse, GetMemoryByIdRequest, GetMemoryByIdResponse,
+        InvalidRequestResponse, KeyDerivationInfo, KeySyncRequest, KeySyncResponse, Memory,
+        MemoryContent, MemoryField, MemoryValue, PlainTextUserInfo, ResetMemoryRequest,
+        ResetMemoryResponse, ResultMask, ScoreRange, SealedMemoryCredentials, SealedMemoryRequest,
+        SealedMemoryResponse, SealedMemorySessionRequest, SealedMemorySessionResponse,
+        SearchMemoryQuery, SearchMemoryRequest, SearchMemoryResponse, SearchMemoryResultItem,
+        UserDb, UserRegistrationRequest, UserRegistrationResponse, WrappedDataEncryptionKey,
         key_sync_response, memory_value, sealed_memory_request, sealed_memory_response,
         search_memory_query, user_registration_response,
     };

oak_private_memory/proto/sealed_memory.proto

@@ -223,6 +223,19 @@ message GetMemoryByIdResponse {
   Memory memory = 2;
 }
 
+message GetMemoriesByIdRequest {
+  repeated string ids = 1;
+  ResultMask result_mask = 2;
+}
+
+message GetMemoriesByIdResponse {
+  // Memories can be returned in an arbitrary order.
+  // If any memory is not found, it will be skipped and the id will be saved
+  // in the `not_found_ids` field.
+  repeated Memory memories = 1;
+  repeated string not_found_ids = 2;
+}
+
 // Metric type for comparing embeddings.
 enum EmbeddingQueryMetricType {
   DOT_PRODUCT = 0;
@@ -405,6 +418,7 @@ message SealedMemoryRequest {
     SearchMemoryRequest search_memory_request = 7;
     UserRegistrationRequest user_registration_request = 8;
     DeleteMemoryRequest delete_memory_request = 9;
+    GetMemoriesByIdRequest get_memories_by_id_request = 10;
   }
 
   // Optional unique identifier for this request within the session.
@@ -424,6 +438,7 @@ message SealedMemoryResponse {
     SearchMemoryResponse search_memory_response = 7;
     UserRegistrationResponse user_registration_response = 8;
     DeleteMemoryResponse delete_memory_response = 9;
+    GetMemoriesByIdResponse get_memories_by_id_response = 10;
 
     // A non-OK status result for a given request.
     //
@@ -439,7 +454,7 @@ message SealedMemoryResponse {
     // TODO: b/474398323 - Update clients to check for this variant.
     // TODO: b/474398548 - Update Sealed Memory server to populate this variant
     // for error conditions.
-    google.rpc.Status error_status = 100;
+    google.rpc.Status error = 100;
   }
 
   // Propagated from the request_id from the request.

oak_private_memory/src/client.rs

@@ -295,4 +295,15 @@ impl PrivateMemoryClient {
             self.invoke(sealed_memory_request::Request::ResetMemoryRequest(request)).await?;
         expect_response_type!(response, sealed_memory_response::Response::ResetMemoryResponse)
     }
+
+    pub async fn get_memories_by_id(
+        &mut self,
+        ids: Vec<String>,
+        result_mask: Option<ResultMask>,
+    ) -> Result<GetMemoriesByIdResponse> {
+        let request = GetMemoriesByIdRequest { ids, result_mask };
+        let response =
+            self.invoke(sealed_memory_request::Request::GetMemoriesByIdRequest(request)).await?;
+        expect_response_type!(response, sealed_memory_response::Response::GetMemoriesByIdResponse)
+    }
 }

oak_private_memory/src/metrics.rs

@@ -394,6 +394,7 @@ impl RequestMetricName {
             sealed_memory_request::Request::GetMemoryByIdRequest(r) => get_name(r),
             sealed_memory_request::Request::SearchMemoryRequest(r) => get_name(r),
             sealed_memory_request::Request::DeleteMemoryRequest(r) => get_name(r),
+            sealed_memory_request::Request::GetMemoriesByIdRequest(r) => get_name(r),
         }))
     }
 }

oak_private_memory/test/client_test.rs

@@ -783,3 +783,93 @@ async fn search_test_embedding_with_expired_memories() {
         assert_eq!(returned_ids.len(), 2);
     }
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_get_memories_by_id() {
+    let (addr, _server_join_handle, _db_join_handle, _persistence_join_handle) =
+        start_server().await.unwrap();
+    let url = format!("http://{}", addr);
+    let pm_uid = "test_get_memories_by_id_user";
+
+    for &format in [SerializationFormat::BinaryProto, SerializationFormat::Json].iter() {
+        let mut client =
+            PrivateMemoryClient::create_with_start_session(&url, pm_uid, TEST_EK, format)
+                .await
+                .unwrap();
+
+        // Add three memories
+        let memory1 = Memory {
+            id: "memory1".to_string(),
+            tags: vec!["tag1".to_string()],
+            ..Default::default()
+        };
+        let memory2 = Memory {
+            id: "memory2".to_string(),
+            tags: vec!["tag2".to_string()],
+            ..Default::default()
+        };
+        let memory3 = Memory {
+            id: "memory3".to_string(),
+            tags: vec!["tag3".to_string()],
+            ..Default::default()
+        };
+
+        client.add_memory(memory1).await.unwrap();
+        client.add_memory(memory2).await.unwrap();
+        client.add_memory(memory3).await.unwrap();
+
+        // Test fetching multiple memories by ID
+        let response = client
+            .get_memories_by_id(
+                vec!["memory3".to_string(), "memory1".to_string(), "memory2".to_string()],
+                None,
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(response.memories.len(), 3);
+        assert!(response.not_found_ids.is_empty());
+        let returned_ids: HashSet<String> =
+            response.memories.iter().map(|m| m.id.clone()).collect();
+        assert!(returned_ids.contains("memory1"));
+        assert!(returned_ids.contains("memory2"));
+        assert!(returned_ids.contains("memory3"));
+
+        // Test fetching a single memory by ID
+        let response = client.get_memories_by_id(vec!["memory2".to_string()], None).await.unwrap();
+        assert_eq!(response.memories.len(), 1);
+        assert_eq!(response.memories[0].id, "memory2");
+        assert!(response.not_found_ids.is_empty());
+
+        // Test fetching with a non-existent ID - should return found ones and report
+        // not found
+        let response = client
+            .get_memories_by_id(
+                vec![
+                    "memory1".to_string(),
+                    "non_existent_id".to_string(),
+                    "memory3".to_string(),
+                    "another_missing".to_string(),
+                ],
+                None,
+            )
+            .await
+            .unwrap();
+        assert_eq!(response.memories.len(), 2);
+        let returned_ids: HashSet<String> =
+            response.memories.iter().map(|m| m.id.clone()).collect();
+        assert!(returned_ids.contains("memory1"));
+        assert!(returned_ids.contains("memory3"));
+        assert_eq!(response.not_found_ids.len(), 2);
+        assert!(response.not_found_ids.contains(&"non_existent_id".to_string()));
+        assert!(response.not_found_ids.contains(&"another_missing".to_string()));
+
+        // Test with all non-existent IDs
+        let response = client
+            .get_memories_by_id(vec!["missing1".to_string(), "missing2".to_string()], None)
+            .await
+            .unwrap();
+        assert!(response.memories.is_empty());
+        assert_eq!(response.not_found_ids.len(), 2);
+    }
+}