Add method to retrieve document by id.

Flag: android.app.appfunctions.flags.enable_app_function_manager
Test: atest FrameworksAppFunctionsTests -c
Bug: 357551503

Change-Id: I801f5c15f248689376dfe405aace8895799b34fa

2 files changed, 90 insertions, 0 deletions

diff --git a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
index 03dd5dd62e51..094723814e17 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
@@ -22,6 +22,9 @@ import android.app.appsearch.AppSearchManager;
 import android.app.appsearch.AppSearchManager.SearchContext;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.AppSearchSession;
+import android.app.appsearch.BatchResultCallback;
+import android.app.appsearch.GenericDocument;
+import android.app.appsearch.GetByDocumentIdRequest;
 import android.app.appsearch.GetSchemaResponse;
 import android.app.appsearch.PutDocumentsRequest;
 import android.app.appsearch.SearchResult;
@@ -189,4 +192,58 @@ public class FutureAppSearchSession implements Closeable {
                     });
         }
     }
+
+    /** A future API to retrieve a document by its id from the local AppSearch session. */
+    public AndroidFuture<GenericDocument> getByDocumentId(
+            @NonNull String documentId, @NonNull String namespace) {
+        Objects.requireNonNull(documentId);
+        Objects.requireNonNull(namespace);
+
+        GetByDocumentIdRequest request =
+                new GetByDocumentIdRequest.Builder(namespace)
+                        .addIds(documentId)
+                        .build();
+        return getSessionAsync()
+                .thenCompose(
+                        session -> {
+                            AndroidFuture<AppSearchBatchResult<String, GenericDocument>>
+                                    batchResultFuture = new AndroidFuture<>();
+                            session.getByDocumentId(
+                                    request,
+                                    mExecutor,
+                                    new BatchResultCallbackAdapter<>(batchResultFuture));
+
+                            return batchResultFuture.thenApply(
+                                    batchResult ->
+                                            getGenericDocumentFromBatchResult(
+                                                    batchResult, documentId));
+                        });
+    }
+
+    private static GenericDocument getGenericDocumentFromBatchResult(
+            AppSearchBatchResult<String, GenericDocument> result, String documentId) {
+        if (result.isSuccess()) {
+            return result.getSuccesses().get(documentId);
+        }
+        throw new IllegalArgumentException("No document in the result for id: " + documentId);
+    }
+
+    private static final class BatchResultCallbackAdapter<K, V>
+            implements BatchResultCallback<K, V> {
+        private final AndroidFuture<AppSearchBatchResult<K, V>> mFuture;
+
+        BatchResultCallbackAdapter(AndroidFuture<AppSearchBatchResult<K, V>> future) {
+            mFuture = future;
+        }
+
+        @Override
+        public void onResult(@NonNull AppSearchBatchResult<K, V> result) {
+            mFuture.complete(result);
+        }
+
+        @Override
+        public void onSystemError(Throwable t) {
+            mFuture.completeExceptionally(t);
+        }
+    }
 }
diff --git a/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt b/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt
index a323799e85e5..a0f1a559bb52 100644
--- a/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt
+++ b/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt
@@ -16,6 +16,7 @@
 package com.android.server.appfunctions
 
 import android.app.appfunctions.AppFunctionRuntimeMetadata
+import android.app.appfunctions.AppFunctionRuntimeMetadata.APP_FUNCTION_RUNTIME_NAMESPACE
 import android.app.appfunctions.AppFunctionRuntimeMetadata.createAppFunctionRuntimeSchema
 import android.app.appfunctions.AppFunctionRuntimeMetadata.createParentAppFunctionRuntimeSchema
 import android.app.appsearch.AppSearchManager
@@ -123,6 +124,38 @@ class FutureAppSearchSessionTest {
         }
     }
 
+    @Test
+    fun getByDocumentId() {
+        val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
+        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+            val setSchemaRequest =
+                SetSchemaRequest.Builder()
+                    .addSchemas(
+                        createParentAppFunctionRuntimeSchema(),
+                        createAppFunctionRuntimeSchema(TEST_PACKAGE_NAME)
+                    )
+                    .build()
+            val schema = session.setSchema(setSchemaRequest)
+            val appFunctionRuntimeMetadata =
+                AppFunctionRuntimeMetadata.Builder(TEST_PACKAGE_NAME, TEST_FUNCTION_ID, "").build()
+            val putDocumentsRequest: PutDocumentsRequest =
+                PutDocumentsRequest.Builder()
+                    .addGenericDocuments(appFunctionRuntimeMetadata)
+                    .build()
+            val putResult = session.put(putDocumentsRequest)
+
+            val genricDocument = session
+                .getByDocumentId(
+                    /* documentId= */ "${TEST_PACKAGE_NAME}/${TEST_FUNCTION_ID}",
+                    APP_FUNCTION_RUNTIME_NAMESPACE
+                )
+                .get()
+
+            val foundAppFunctionRuntimeMetadata = AppFunctionRuntimeMetadata(genricDocument)
+            assertThat(foundAppFunctionRuntimeMetadata.functionId).isEqualTo(TEST_FUNCTION_ID)
+        }
+    }
+
     private companion object {
         const val TEST_DB: String = "test_db"
         const val TEST_PACKAGE_NAME: String = "test_pkg"