Support queries in AppSearchImpl and FakeIcing.

Bug: 145631811
Test: atest CtsAppSearchTestCases FrameworksCoreTests:android.app.appsearch FrameworksServicesTests:com.android.server.appsearch.impl
Change-Id: Ic47da9bca664999c5ba679a81e0c7e9d7471d4de
diff --git a/framework/java/android/app/appsearch/AppSearchManager.java b/framework/java/android/app/appsearch/AppSearchManager.java
index e5a1639..1fb7b4b 100644
--- a/framework/java/android/app/appsearch/AppSearchManager.java
+++ b/framework/java/android/app/appsearch/AppSearchManager.java
@@ -15,7 +15,6 @@
  */
 package android.app.appsearch;
 
-import android.annotation.CallbackExecutor;
 import android.annotation.NonNull;
 import android.annotation.SystemService;
 import android.content.Context;
@@ -35,8 +34,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Executor;
-import java.util.function.BiConsumer;
 
 /**
  * This class provides access to the centralized AppSearch index maintained by the system.
@@ -82,8 +79,8 @@
      *     <li>Removal of an existing type
      *     <li>Removal of a property from a type
      *     <li>Changing the data type ({@code boolean}, {@code long}, etc.) of an existing property
-     *     <li>For properties of {@code Document} type, changing the schema type of
-     *         {@code Document Documents} of that property
+     *     <li>For properties of {@code AppSearchDocument} type, changing the schema type of
+     *         {@code AppSearchDocument}s of that property
      *     <li>Changing the cardinality of a data type to be more restrictive (e.g. changing an
      *         {@link android.app.appsearch.AppSearchSchema.PropertyConfig#CARDINALITY_OPTIONAL
      *             OPTIONAL} property into a
@@ -156,15 +153,15 @@
     }
 
     /**
-     * Index {@link AppSearchDocument Documents} into AppSearch.
+     * Index {@link AppSearchDocument}s into AppSearch.
      *
      * <p>You should not call this method directly; instead, use the
      * {@code AppSearch#putDocuments()} API provided by JetPack.
      *
-     * <p>Each {@link AppSearchDocument Document's} {@code schemaType} field must be set to the
-     * name of a schema type previously registered via the {@link #setSchema} method.
+     * <p>Each {@link AppSearchDocument}'s {@code schemaType} field must be set to the name of a
+     * schema type previously registered via the {@link #setSchema} method.
      *
-     * @param documents {@link AppSearchDocument Documents} that need to be indexed.
+     * @param documents {@link AppSearchDocument}s that need to be indexed.
      * @return An {@link AppSearchBatchResult} mapping the document URIs to {@link Void} if they
      *     were successfully indexed, or a {@link Throwable} describing the failure if they could
      *     not be indexed.
@@ -253,10 +250,12 @@
     }
 
     /**
-     * This method searches for documents based on a given query string. It also accepts
-     * specifications regarding how to search and format the results.
+     * Searches a document based on a given query string.
      *
-     *<p>Currently we support following features in the raw query format:
+     * <p>You should not call this method directly; instead, use the {@code AppSearch#query()} API
+     * provided by JetPack.
+     *
+     * <p>Currently we support following features in the raw query format:
      * <ul>
      *     <li>AND
      *     <p>AND joins (e.g. “match documents that have both the terms ‘dog’ and
@@ -288,59 +287,50 @@
      *     ‘Video’ schema type.
      * </ul>
      *
-     * <p> It is strongly recommended to use Jetpack APIs.
-     *
      * @param queryExpression Query String to search.
      * @param searchSpec Spec for setting filters, raw query etc.
-     * @param executor Executor on which to invoke the callback.
-     * @param callback  Callback to receive errors resulting from the query operation. If the
-     *                 operation succeeds, the callback will be invoked with {@code null}.
      * @hide
      */
     @NonNull
-    public void query(
-            @NonNull String queryExpression,
-            @NonNull SearchSpec searchSpec,
-            @NonNull @CallbackExecutor Executor executor,
-            @NonNull BiConsumer<? super SearchResults, ? super Throwable> callback) {
-        AndroidFuture<byte[]> future = new AndroidFuture<>();
-        future.whenCompleteAsync((searchResultBytes, err) -> {
-            if (err != null) {
-                callback.accept(null, err);
-                return;
-            }
-            if (searchResultBytes != null) {
-                SearchResultProto searchResultProto;
-                try {
-                    searchResultProto = SearchResultProto.parseFrom(searchResultBytes);
-                } catch (InvalidProtocolBufferException e) {
-                    callback.accept(null, e);
-                    return;
-                }
-                if (searchResultProto.getStatus().getCode() != StatusProto.Code.OK) {
-                    // TODO(sidchhabra): Add better exception handling.
-                    callback.accept(
-                            null,
-                            new RuntimeException(searchResultProto.getStatus().getMessage()));
-                    return;
-                }
-                SearchResults searchResults = new SearchResults(searchResultProto);
-                callback.accept(searchResults, null);
-                return;
-            }
-            // Nothing was supplied in the future at all
-            callback.accept(
-                    null, new IllegalStateException("Unknown failure occurred while querying"));
-        }, executor);
+    public AppSearchResult<SearchResults> query(
+            @NonNull String queryExpression, @NonNull SearchSpec searchSpec) {
+        // TODO(b/146386470): Transmit the result documents as a RemoteStream instead of sending
+        //     them in one big list.
+        AndroidFuture<AppSearchResult> searchResultFuture = new AndroidFuture<>();
         try {
             SearchSpecProto searchSpecProto = searchSpec.getSearchSpecProto();
             searchSpecProto = searchSpecProto.toBuilder().setQuery(queryExpression).build();
-            mService.query(searchSpecProto.toByteArray(),
+            mService.query(
+                    searchSpecProto.toByteArray(),
                     searchSpec.getResultSpecProto().toByteArray(),
-                    searchSpec.getScoringSpecProto().toByteArray(), future);
+                    searchSpec.getScoringSpecProto().toByteArray(),
+                    searchResultFuture);
         } catch (RemoteException e) {
-            future.completeExceptionally(e);
+            searchResultFuture.completeExceptionally(e);
         }
+
+        // Deserialize the protos into Document objects
+        AppSearchResult<byte[]> searchResultBytes = getFutureOrThrow(searchResultFuture);
+        if (!searchResultBytes.isSuccess()) {
+            return AppSearchResult.newFailedResult(
+                    searchResultBytes.getResultCode(), searchResultBytes.getErrorMessage());
+        }
+        SearchResultProto searchResultProto;
+        try {
+            searchResultProto = SearchResultProto.parseFrom(searchResultBytes.getResultValue());
+        } catch (InvalidProtocolBufferException e) {
+            return AppSearchResult.newFailedResult(
+                    AppSearchResult.RESULT_INTERNAL_ERROR, e.getMessage());
+        }
+        if (searchResultProto.getStatus().getCode() != StatusProto.Code.OK) {
+            // This should never happen; AppSearchManagerService should catch failed searchResults
+            // entries and transmit them as a failed AppSearchResult.
+            return AppSearchResult.newFailedResult(
+                    AppSearchResult.RESULT_INTERNAL_ERROR,
+                    searchResultProto.getStatus().getMessage());
+        }
+
+        return AppSearchResult.newSuccessfulResult(new SearchResults(searchResultProto));
     }
 
     private static <T> T getFutureOrThrow(@NonNull AndroidFuture<T> future) {
diff --git a/framework/java/android/app/appsearch/IAppSearchManager.aidl b/framework/java/android/app/appsearch/IAppSearchManager.aidl
index 8e3f1e7..b7100a4 100644
--- a/framework/java/android/app/appsearch/IAppSearchManager.aidl
+++ b/framework/java/android/app/appsearch/IAppSearchManager.aidl
@@ -66,9 +66,10 @@
      * @param searchSpecBytes Serialized SearchSpecProto.
      * @param resultSpecBytes Serialized SearchResultsProto.
      * @param scoringSpecBytes Serialized ScoringSpecProto.
-     * @param callback {@link AndroidFuture}. Will be completed with a serialized
-     *     {@link SearchResultsProto}, or completed exceptionally if query fails.
+     * @param callback {@link AndroidFuture}&lt;{@link AppSearchResult}&lt;{@link byte[]}&gt;&gt;
+     *     Will be completed with a serialized {@link SearchResultsProto}.
      */
-    void query(in byte[] searchSpecBytes, in byte[] resultSpecBytes,
-            in byte[] scoringSpecBytes, in AndroidFuture callback);
+    void query(
+        in byte[] searchSpecBytes, in byte[] resultSpecBytes, in byte[] scoringSpecBytes,
+        in AndroidFuture<AppSearchResult> callback);
 }
diff --git a/service/java/com/android/server/appsearch/AppSearchManagerService.java b/service/java/com/android/server/appsearch/AppSearchManagerService.java
index d26b4e3..5cc5add 100644
--- a/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -27,14 +27,15 @@
 import com.android.internal.util.Preconditions;
 import com.android.server.SystemService;
 import com.android.server.appsearch.impl.AppSearchImpl;
-import com.android.server.appsearch.impl.FakeIcing;
 import com.android.server.appsearch.impl.ImplInstanceManager;
 
 import com.google.android.icing.proto.DocumentProto;
+import com.google.android.icing.proto.ResultSpecProto;
 import com.google.android.icing.proto.SchemaProto;
+import com.google.android.icing.proto.ScoringSpecProto;
 import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
-import com.google.android.icing.protobuf.InvalidProtocolBufferException;
+import com.google.android.icing.proto.StatusProto;
 
 import java.io.IOException;
 import java.util.List;
@@ -46,11 +47,8 @@
 
     public AppSearchManagerService(Context context) {
         super(context);
-        mFakeIcing = new FakeIcing();
     }
 
-    private final FakeIcing mFakeIcing;
-
     @Override
     public void onStart() {
         publishBinderService(Context.APP_SEARCH_SERVICE, new Stub());
@@ -144,23 +142,43 @@
             }
         }
 
-        // TODO(sidchhabra):Init FakeIcing properly.
         // TODO(sidchhabra): Do this in a threadpool.
         @Override
-        public void query(@NonNull byte[] searchSpec, @NonNull byte[] resultSpec,
-                @NonNull byte[] scoringSpec, AndroidFuture callback) {
-            Preconditions.checkNotNull(searchSpec);
-            Preconditions.checkNotNull(resultSpec);
-            Preconditions.checkNotNull(scoringSpec);
-            SearchSpecProto searchSpecProto = null;
+        public void query(
+                @NonNull byte[] searchSpecBytes,
+                @NonNull byte[] resultSpecBytes,
+                @NonNull byte[] scoringSpecBytes,
+                @NonNull AndroidFuture<AppSearchResult> callback) {
+            Preconditions.checkNotNull(searchSpecBytes);
+            Preconditions.checkNotNull(resultSpecBytes);
+            Preconditions.checkNotNull(scoringSpecBytes);
+            Preconditions.checkNotNull(callback);
+            int callingUid = Binder.getCallingUidOrThrow();
+            int callingUserId = UserHandle.getUserId(callingUid);
+            long callingIdentity = Binder.clearCallingIdentity();
             try {
-                searchSpecProto = SearchSpecProto.parseFrom(searchSpec);
-            } catch (InvalidProtocolBufferException e) {
-                throw new RuntimeException(e);
+                SearchSpecProto searchSpecProto = SearchSpecProto.parseFrom(searchSpecBytes);
+                ResultSpecProto resultSpecProto = ResultSpecProto.parseFrom(resultSpecBytes);
+                ScoringSpecProto scoringSpecProto = ScoringSpecProto.parseFrom(scoringSpecBytes);
+                AppSearchImpl impl = ImplInstanceManager.getInstance(getContext(), callingUserId);
+                SearchResultProto searchResultProto =
+                        impl.query(callingUid, searchSpecProto, resultSpecProto, scoringSpecProto);
+                // TODO(sidchhabra): Translate SearchResultProto errors into error codes. This might
+                //     better be done in AppSearchImpl by throwing an AppSearchException.
+                if (searchResultProto.getStatus().getCode() != StatusProto.Code.OK) {
+                    callback.complete(
+                            AppSearchResult.newFailedResult(
+                                    AppSearchResult.RESULT_INTERNAL_ERROR,
+                                    searchResultProto.getStatus().getMessage()));
+                } else {
+                    callback.complete(
+                            AppSearchResult.newSuccessfulResult(searchResultProto.toByteArray()));
+                }
+            } catch (Throwable t) {
+                callback.complete(throwableToFailedResult(t));
+            } finally {
+                Binder.restoreCallingIdentity(callingIdentity);
             }
-            SearchResultProto searchResults =
-                    mFakeIcing.query(searchSpecProto.getQuery());
-            callback.complete(searchResults.toByteArray());
         }
 
         private <ValueType> AppSearchResult<ValueType> throwableToFailedResult(
diff --git a/service/java/com/android/server/appsearch/impl/AppSearchImpl.java b/service/java/com/android/server/appsearch/impl/AppSearchImpl.java
index 7442d06..f0af398 100644
--- a/service/java/com/android/server/appsearch/impl/AppSearchImpl.java
+++ b/service/java/com/android/server/appsearch/impl/AppSearchImpl.java
@@ -20,14 +20,21 @@
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.content.Context;
+import android.util.ArraySet;
 
 import com.android.internal.annotations.VisibleForTesting;
 
 import com.google.android.icing.proto.DocumentProto;
 import com.google.android.icing.proto.PropertyConfigProto;
 import com.google.android.icing.proto.PropertyProto;
+import com.google.android.icing.proto.ResultSpecProto;
 import com.google.android.icing.proto.SchemaProto;
 import com.google.android.icing.proto.SchemaTypeConfigProto;
+import com.google.android.icing.proto.ScoringSpecProto;
+import com.google.android.icing.proto.SearchResultProto;
+import com.google.android.icing.proto.SearchSpecProto;
+
+import java.util.Set;
 
 /**
  * Manages interaction with {@link FakeIcing} and other components to implement AppSearch
@@ -122,6 +129,14 @@
     public DocumentProto getDocument(int callingUid, @NonNull String uri) {
         String typePrefix = getTypePrefix(callingUid);
         DocumentProto document = mFakeIcing.get(uri);
+
+        // TODO(b/146526096): Since FakeIcing doesn't currently handle namespaces, we perform a
+        //  post-filter to make sure we don't return documents we shouldn't. This should be removed
+        //  once the real Icing Lib is implemented.
+        if (!document.getNamespace().equals(typePrefix)) {
+            return null;
+        }
+
         // Rewrite the type names to remove the app's prefix
         DocumentProto.Builder documentBuilder = document.toBuilder();
         rewriteDocumentTypes(typePrefix, documentBuilder, /*add=*/ false);
@@ -129,6 +144,71 @@
     }
 
     /**
+     * Executes a query against the AppSearch index and returns results.
+     *
+     * @param callingUid The uid of the app calling AppSearch.
+     * @param searchSpec Defines what and how to search
+     * @param resultSpec Defines what results to show
+     * @param scoringSpec Defines how to order results
+     * @return The results of performing this search  The proto might have no {@code results} if no
+     *     documents matched the query.
+     */
+    @NonNull
+    public SearchResultProto query(
+            int callingUid,
+            @NonNull SearchSpecProto searchSpec,
+            @NonNull ResultSpecProto resultSpec,
+            @NonNull ScoringSpecProto scoringSpec) {
+        String typePrefix = getTypePrefix(callingUid);
+        SearchResultProto searchResults = mFakeIcing.query(searchSpec.getQuery());
+        if (searchResults.getResultsCount() == 0) {
+            return searchResults;
+        }
+        Set<String> qualifiedSearchFilters = null;
+        if (searchSpec.getSchemaTypeFiltersCount() > 0) {
+            qualifiedSearchFilters = new ArraySet<>(searchSpec.getSchemaTypeFiltersCount());
+            for (String schema : searchSpec.getSchemaTypeFiltersList()) {
+                String qualifiedSchema = typePrefix + schema;
+                qualifiedSearchFilters.add(qualifiedSchema);
+            }
+        }
+        // Rewrite the type names to remove the app's prefix
+        SearchResultProto.Builder searchResultsBuilder = searchResults.toBuilder();
+        for (int i = 0; i < searchResultsBuilder.getResultsCount(); i++) {
+            if (searchResults.getResults(i).hasDocument()) {
+                SearchResultProto.ResultProto.Builder resultBuilder =
+                        searchResultsBuilder.getResults(i).toBuilder();
+
+                // TODO(b/145631811): Since FakeIcing doesn't currently handle namespaces, we
+                //  perform a post-filter to make sure we don't return documents we shouldn't. This
+                //  should be removed once the real Icing Lib is implemented.
+                if (!resultBuilder.getDocument().getNamespace().equals(typePrefix)) {
+                    searchResultsBuilder.removeResults(i);
+                    i--;
+                    continue;
+                }
+
+                // TODO(b/145631811): Since FakeIcing doesn't currently handle type names, we
+                //  perform a post-filter to make sure we don't return documents we shouldn't. This
+                //  should be removed once the real Icing Lib is implemented.
+                if (qualifiedSearchFilters != null
+                        && !qualifiedSearchFilters.contains(
+                                resultBuilder.getDocument().getSchema())) {
+                    searchResultsBuilder.removeResults(i);
+                    i--;
+                    continue;
+                }
+
+                DocumentProto.Builder documentBuilder = resultBuilder.getDocument().toBuilder();
+                rewriteDocumentTypes(typePrefix, documentBuilder, /*add=*/false);
+                resultBuilder.setDocument(documentBuilder);
+                searchResultsBuilder.setResults(i, resultBuilder);
+            }
+        }
+        return searchResultsBuilder.build();
+    }
+
+    /**
      * Rewrites all types mentioned anywhere in {@code documentBuilder} to prepend or remove
      * {@code typePrefix}.
      *
diff --git a/service/java/com/android/server/appsearch/impl/FakeIcing.java b/service/java/com/android/server/appsearch/impl/FakeIcing.java
index d07ef4b..fce0c80 100644
--- a/service/java/com/android/server/appsearch/impl/FakeIcing.java
+++ b/service/java/com/android/server/appsearch/impl/FakeIcing.java
@@ -88,22 +88,35 @@
     }
 
     /**
-     * Returns documents containing the given term.
+     * Returns documents containing all words in the given query string.
      *
-     * @param term A single exact term to look up in the index.
+     * @param queryExpression A set of words to search for. They will be implicitly AND-ed together.
+     *     No operators are supported.
      * @return A {@link SearchResultProto} containing the matching documents, which may have no
      *   results if no documents match.
      */
     @NonNull
-    public SearchResultProto query(@NonNull String term) {
-        String normTerm = normalizeString(term);
-        Set<Integer> docIds = mIndex.get(normTerm);
+    public SearchResultProto query(@NonNull String queryExpression) {
+        String[] terms = normalizeString(queryExpression).split("\\s+");
         SearchResultProto.Builder results = SearchResultProto.newBuilder()
                 .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.OK));
+        if (terms.length == 0) {
+            return results.build();
+        }
+        Set<Integer> docIds = mIndex.get(terms[0]);
         if (docIds == null || docIds.isEmpty()) {
             return results.build();
         }
-
+        for (int i = 1; i < terms.length; i++) {
+            Set<Integer> termDocIds = mIndex.get(terms[i]);
+            if (termDocIds == null) {
+                return results.build();
+            }
+            docIds.retainAll(termDocIds);
+            if (docIds.isEmpty()) {
+                return results.build();
+            }
+        }
         for (int docId : docIds) {
             DocumentProto document = mDocStore.get(docId);
             if (document != null) {
diff --git a/testing/coretests/src/android/app/appsearch/SearchResultsTest.java b/testing/coretests/src/android/app/appsearch/SearchResultsTest.java
index 21259cc..67cc53c 100644
--- a/testing/coretests/src/android/app/appsearch/SearchResultsTest.java
+++ b/testing/coretests/src/android/app/appsearch/SearchResultsTest.java
@@ -16,7 +16,7 @@
 
 package android.app.appsearch;
 
-import static  com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.assertThrows;