Merge "Implement getDocuments() support in AppSearchImpl."

4 files changed, 137 insertions, 12 deletions

diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchManager.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchManager.java
index e2c9b0f74870..39abe7813578 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchManager.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchManager.java
@@ -23,6 +23,7 @@ import android.os.RemoteException;
 
 import com.android.internal.infra.AndroidFuture;
 
+import com.google.android.icing.proto.DocumentProto;
 import com.google.android.icing.proto.SchemaProto;
 import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
@@ -181,6 +182,53 @@ public class AppSearchManager {
     }
 
     /**
+     * Retrieves {@link android.app.appsearch.AppSearch.Document}s by URI.
+     *
+     * <p>You should not call this method directly; instead, use the
+     * {@code AppSearch#getDocuments()} API provided by JetPack.
+     *
+     * @param uris URIs of the documents to look up.
+     * @param executor Executor on which to invoke the callback.
+     * @param callback Callback to receive the documents or error.
+     */
+    public void getDocuments(
+            @NonNull List<String> uris,
+            @NonNull @CallbackExecutor Executor executor,
+            @NonNull BiConsumer<List<AppSearch.Document>, ? super Throwable> callback) {
+        AndroidFuture<List<byte[]>> future = new AndroidFuture<>();
+        future.whenCompleteAsync((documentProtos, err) -> {
+            if (err != null) {
+                callback.accept(null, err);
+                return;
+            }
+            if (documentProtos != null) {
+                List<AppSearch.Document> results = new ArrayList<>(documentProtos.size());
+                for (int i = 0; i < documentProtos.size(); i++) {
+                    DocumentProto documentProto;
+                    try {
+                        documentProto = DocumentProto.parseFrom(documentProtos.get(i));
+                    } catch (InvalidProtocolBufferException e) {
+                        callback.accept(null, e);
+                        return;
+                    }
+                    results.add(new AppSearch.Document(documentProto));
+                }
+                callback.accept(results, null);
+                return;
+            }
+            // Nothing was supplied in the future at all
+            callback.accept(null, new IllegalStateException(
+                    "Unknown failure occurred while retrieving documents"));
+        }, executor);
+        // TODO(b/146386470) stream uris?
+        try {
+            mService.getDocuments(uris.toArray(new String[uris.size()]), future);
+        } catch (RemoteException e) {
+            future.completeExceptionally(e);
+        }
+    }
+
+    /**
      * This method searches for documents based on a given query string. It also accepts
      * specifications regarding how to search and format the results.
      *
@@ -237,7 +285,6 @@ public class AppSearchManager {
                 callback.accept(null, err);
                 return;
             }
-
             if (searchResultBytes != null) {
                 SearchResultProto searchResultProto;
                 try {
@@ -257,12 +304,10 @@ public class AppSearchManager {
                 callback.accept(searchResults, null);
                 return;
             }
-
             // Nothing was supplied in the future at all
             callback.accept(
                     null, new IllegalStateException("Unknown failure occurred while querying"));
         }, executor);
-
         try {
             SearchSpecProto searchSpecProto = searchSpec.getSearchSpecProto();
             searchSpecProto = searchSpecProto.toBuilder().setQuery(queryExpression).build();
diff --git a/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl b/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
index eef41ed7104d..b80082536ebd 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
+++ b/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
@@ -47,6 +47,16 @@ interface IAppSearchManager {
     void putDocuments(in List documentsBytes, in AndroidFuture<AppSearchBatchResult> callback);
 
     /**
+     * Retrieves documents from the index.
+     *
+     * @param uris The URIs of the documents to retrieve
+     * @param callback {@link AndroidFuture}&lt;{@link List}&lt;byte[]&gt;&gt;. Will be completed
+     *     with a {@link List} containing serialized DocumentProtos, or completed exceptionally if
+     *     get fails.
+     */
+    void getDocuments(in String[] uris, in AndroidFuture callback);
+
+    /**
      * Searches a document based on a given specifications.
      *
      * @param searchSpecBytes Serialized SearchSpecProto.
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
index 6293ee7059e5..2608e8453b3e 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -35,6 +35,7 @@ import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
 import com.google.android.icing.protobuf.InvalidProtocolBufferException;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -103,6 +104,30 @@ public class AppSearchManagerService extends SystemService {
                 Binder.restoreCallingIdentity(callingIdentity);
             }
         }
+
+        @Override
+        public void getDocuments(String[] uris, AndroidFuture callback) {
+            Preconditions.checkNotNull(uris);
+            Preconditions.checkNotNull(callback);
+            int callingUid = Binder.getCallingUidOrThrow();
+            int callingUserId = UserHandle.getUserId(callingUid);
+            long callingIdentity = Binder.clearCallingIdentity();
+            try {
+                AppSearchImpl impl = ImplInstanceManager.getInstance(getContext(), callingUserId);
+                // Contains serialized DocumentProto. byte[][] is not transmissible via Binder.
+                List<byte[]> results = new ArrayList<>(uris.length);
+                for (String uri : uris) {
+                    DocumentProto result = impl.getDocument(callingUid, uri);
+                    results.add(result.toByteArray());
+                }
+                callback.complete(results);
+            } catch (Throwable t) {
+                callback.completeExceptionally(t);
+            } finally {
+                Binder.restoreCallingIdentity(callingIdentity);
+            }
+        }
+
         // TODO(sidchhabra):Init FakeIcing properly.
         // TODO(sidchhabra): Do this in a threadpool.
         @Override
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/impl/AppSearchImpl.java b/apex/appsearch/service/java/com/android/server/appsearch/impl/AppSearchImpl.java
index 04b4b1427328..7442d0625280 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/impl/AppSearchImpl.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/impl/AppSearchImpl.java
@@ -17,6 +17,7 @@
 package com.android.server.appsearch.impl;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.content.Context;
 
@@ -106,30 +107,65 @@ public final class AppSearchImpl {
         // Rewrite the type names to include the app's prefix
         String typePrefix = getTypePrefix(callingUid);
         DocumentProto.Builder documentBuilder = origDocument.toBuilder();
-        rewriteDocumentTypes(typePrefix, documentBuilder);
+        rewriteDocumentTypes(typePrefix, documentBuilder, /*add=*/ true);
         mFakeIcing.put(documentBuilder.build());
     }
 
     /**
-     * Rewrites all types mentioned anywhere in {@code documentBuilder} to prepend
+     * Retrieves a document from the AppSearch index by URI.
+     *
+     * @param callingUid The uid of the app calling AppSearch.
+     * @param uri The URI of the document to get.
+     * @return The Document contents, or {@code null} if no such URI exists in the system.
+     */
+    @Nullable
+    public DocumentProto getDocument(int callingUid, @NonNull String uri) {
+        String typePrefix = getTypePrefix(callingUid);
+        DocumentProto document = mFakeIcing.get(uri);
+        // Rewrite the type names to remove the app's prefix
+        DocumentProto.Builder documentBuilder = document.toBuilder();
+        rewriteDocumentTypes(typePrefix, documentBuilder, /*add=*/ false);
+        return documentBuilder.build();
+    }
+
+    /**
+     * Rewrites all types mentioned anywhere in {@code documentBuilder} to prepend or remove
      * {@code typePrefix}.
      *
-     * @param typePrefix The prefix to add
+     * @param typePrefix The prefix to add or remove
      * @param documentBuilder The document to mutate
+     * @param add Whether to add typePrefix to the types. If {@code false}, typePrefix will be
+     *     removed from the types.
+     * @throws IllegalArgumentException If {@code add=false} and the document has a type that
+     *     doesn't start with {@code typePrefix}.
      */
     @VisibleForTesting
     void rewriteDocumentTypes(
             @NonNull String typePrefix,
-            @NonNull DocumentProto.Builder documentBuilder) {
-        // Rewrite the type name to include the app's prefix
-        String newSchema = typePrefix + documentBuilder.getSchema();
+            @NonNull DocumentProto.Builder documentBuilder,
+            boolean add) {
+        // Rewrite the type name to include/remove the app's prefix
+        String newSchema;
+        if (add) {
+            newSchema = typePrefix + documentBuilder.getSchema();
+        } else {
+            newSchema = removePrefix(typePrefix, documentBuilder.getSchema());
+        }
         documentBuilder.setSchema(newSchema);
 
-        // Add namespace. If we ever allow users to set their own namespaces, this will have
+        // Add/remove namespace. If we ever allow users to set their own namespaces, this will have
         // to change to prepend the prefix instead of setting the whole namespace. We will also have
         // to store the namespaces in a map similar to the type map so we can rewrite queries with
         // empty namespaces.
-        documentBuilder.setNamespace(typePrefix);
+        if (add) {
+            documentBuilder.setNamespace(typePrefix);
+        } else if (!documentBuilder.getNamespace().equals(typePrefix)) {
+            throw new IllegalStateException(
+                    "Unexpected namespace \"" + documentBuilder.getNamespace()
+                            + "\" (expected \"" + typePrefix + "\")");
+        } else {
+            documentBuilder.clearNamespace();
+        }
 
         // Recurse into derived documents
         for (int propertyIdx = 0;
@@ -142,7 +178,7 @@ public final class AppSearchImpl {
                 for (int documentIdx = 0; documentIdx < documentCount; documentIdx++) {
                     DocumentProto.Builder derivedDocumentBuilder =
                             propertyBuilder.getDocumentValues(documentIdx).toBuilder();
-                    rewriteDocumentTypes(typePrefix, derivedDocumentBuilder);
+                    rewriteDocumentTypes(typePrefix, derivedDocumentBuilder, add);
                     propertyBuilder.setDocumentValues(documentIdx, derivedDocumentBuilder);
                 }
                 documentBuilder.setProperties(propertyIdx, propertyBuilder);
@@ -165,4 +201,13 @@ public final class AppSearchImpl {
         }
         return callingUidName + "@" + mUserId + "/";
     }
+
+    @NonNull
+    private static String removePrefix(@NonNull String prefix, @NonNull String input) {
+        if (!input.startsWith(prefix)) {
+            throw new IllegalArgumentException(
+                    "Input \"" + input + "\" does not start with \"" + prefix + "\"");
+        }
+        return input.substring(prefix.length());
+    }
 }