Fix startAfter/endBefore for orderByKeys queries

firebase-database/src/androidTest/java/com/google/firebase/database/QueryTest.java

@@ -14,6 +14,7 @@
 
 package com.google.firebase.database;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -581,6 +582,111 @@ public void setVariousLimitsWithStartAtName() throws DatabaseException, Interrup
     expectations.waitForEvents();
   }
 
+  @Test
+  public void testStartAfterWithOrderByKey()
+      throws DatabaseException, InterruptedException, ExecutionException, TimeoutException,
+          TestFailure {
+    DatabaseReference ref = IntegrationTestHelpers.getRandomNode();
+    DatabaseReference childOne = ref.push();
+    DatabaseReference childTwo = ref.push();
+    Tasks.await(childOne.setValue(1L));
+    Tasks.await(childTwo.setValue(2L));
+
+    DataSnapshot snapshot = Tasks.await(ref.orderByKey().startAfter(childOne.getKey()).get());
+
+    Map<String, Long> values = (Map<String, Long>) snapshot.getValue();
+
+    assertNotNull(values);
+    assertArrayEquals(values.keySet().toArray(), new String[] {childTwo.getKey()});
+    assertArrayEquals(values.values().toArray(), new Long[] {values.get(childTwo.getKey())});
+  }
+
+  @Test
+  public void testEndBeforeWithOrderByKey()
+      throws DatabaseException, InterruptedException, ExecutionException {
+    DatabaseReference ref = IntegrationTestHelpers.getRandomNode();
+    DatabaseReference childOne = ref.push();
+    DatabaseReference childTwo = ref.push();
+    Tasks.await(childOne.setValue(1L));
+    Tasks.await(childTwo.setValue(2L));
+
+    DataSnapshot snapshot = Tasks.await(ref.orderByKey().endBefore(childTwo.getKey()).get());
+    Map<String, Long> values = (Map<String, Long>) snapshot.getValue();
+
+    assertNotNull(values);
+    assertArrayEquals(values.keySet().toArray(), new String[] {childOne.getKey()});
+    assertArrayEquals(values.values().toArray(), new Long[] {values.get(childOne.getKey())});
+  }
+
+  // This test checks that range filters are applied to in-memory data if our active listeners
+  // have already retrieved the data we need to satisfy the get().
+  @Test
+  public void testEndBeforeWithOrderByKeyOverlappingListener()
+      throws DatabaseException, InterruptedException, ExecutionException {
+    DatabaseReference ref = IntegrationTestHelpers.getRandomNode();
+    DatabaseReference childOne = ref.push();
+    DatabaseReference childTwo = ref.push();
+    Tasks.await(childOne.setValue(1L));
+    Tasks.await(childTwo.setValue(2L));
+
+    Semaphore semaphore = new Semaphore(0);
+    ValueEventListener listener =
+        new ValueEventListener() {
+          @Override
+          public void onDataChange(@NonNull DataSnapshot snapshot) {
+            semaphore.release();
+          }
+
+          @Override
+          public void onCancelled(@NonNull DatabaseError error) {}
+        };
+
+    ref.addValueEventListener(listener);
+
+    IntegrationTestHelpers.waitFor(semaphore);
+
+    DataSnapshot snapshot = Tasks.await(ref.orderByKey().endBefore(childTwo.getKey()).get());
+    Map<String, Long> values = (Map<String, Long>) snapshot.getValue();
+
+    assertNotNull(values);
+    assertArrayEquals(values.keySet().toArray(), new String[] {childOne.getKey()});
+    assertArrayEquals(values.values().toArray(), new Long[] {values.get(childOne.getKey())});
+    ref.removeEventListener(listener);
+  }
+
+  @Test
+  public void testStartAfterWithOrderByKeyOverlappingListener()
+      throws DatabaseException, InterruptedException, ExecutionException {
+    DatabaseReference ref = IntegrationTestHelpers.getRandomNode();
+    DatabaseReference childOne = ref.push();
+    DatabaseReference childTwo = ref.push();
+    Tasks.await(childOne.setValue(1L));
+    Tasks.await(childTwo.setValue(2L));
+
+    Semaphore semaphore = new Semaphore(0);
+    ValueEventListener listener =
+        new ValueEventListener() {
+          @Override
+          public void onDataChange(@NonNull DataSnapshot snapshot) {
+            semaphore.release();
+          }
+
+          @Override
+          public void onCancelled(@NonNull DatabaseError error) {}
+        };
+    ref.addValueEventListener(listener);
+
+    IntegrationTestHelpers.waitFor(semaphore);
+
+    DataSnapshot snapshot = Tasks.await(ref.orderByKey().startAfter(childOne.getKey()).get());
+    Map<String, Long> values = (Map<String, Long>) snapshot.getValue();
+
+    assertNotNull(values);
+    assertArrayEquals(values.keySet().toArray(), new String[] {childTwo.getKey()});
+    assertArrayEquals(values.values().toArray(), new Long[] {values.get(childTwo.getKey())});
+    ref.removeEventListener(listener);
+  }
+
   @Test
   public void setVariousLimitsWithStartAfterName() throws DatabaseException, InterruptedException {
     DatabaseReference ref = IntegrationTestHelpers.getRandomNode();
@@ -4526,6 +4632,33 @@ public void onCancelled(@NonNull DatabaseError error) {}
     }
   }
 
+  @Test
+  public void testGetWithPendingWrites() throws ExecutionException, InterruptedException {
+    DatabaseReference node = IntegrationTestHelpers.getRandomNode();
+    node.getDatabase().goOffline();
+    try {
+      Map<String, Object> expected = new MapBuilder().put("foo", "bar").build();
+      node.setValue(expected);
+      DataSnapshot snapshot = Tasks.await(node.get());
+      assertEquals(snapshot.getValue(), expected);
+    } finally {
+      node.getDatabase().goOnline();
+    }
+  }
+
+  @Test
+  public void testGetChildOfPendingWrites() throws ExecutionException, InterruptedException {
+    DatabaseReference node = IntegrationTestHelpers.getRandomNode();
+    node.getDatabase().goOffline();
+    try {
+      node.setValue(new MapBuilder().put("foo", "bar").build());
+      DataSnapshot snapshot = Tasks.await(node.child("foo").get());
+      assertEquals(snapshot.getValue(), "bar");
+    } finally {
+      node.getDatabase().goOnline();
+    }
+  }
+
   @Test
   public void testGetSendsServerProbesPersistenceCacheWhenOfflineWithNoListener()
       throws DatabaseException, InterruptedException, ExecutionException, TestFailure,

firebase-database/src/main/java/com/google/firebase/database/Query.java

@@ -83,9 +83,10 @@ public class Query {
   private void validateQueryEndpoints(QueryParams params) {
     if (params.getIndex().equals(KeyIndex.getInstance())) {
       String message =
-          "You must use startAt(String value), endAt(String value) or "
-              + "equalTo(String value) in combination with orderByKey(). Other type of values or "
-              + "using the version with 2 parameters is not supported";
+          "You must use startAt(String value), startAfter(String value), endAt(String value), "
+              + "endBefore(String value) or equalTo(String value) in combination with "
+              + "orderByKey(). Other type of values or using the version with 2 parameters is "
+              + "not supported";
       if (params.hasStart()) {
         Node startNode = params.getIndexStartValue();
         ChildKey startName = params.getIndexStartName();
@@ -105,8 +106,8 @@ private void validateQueryEndpoints(QueryParams params) {
       if ((params.hasStart() && !PriorityUtilities.isValidPriority(params.getIndexStartValue()))
           || (params.hasEnd() && !PriorityUtilities.isValidPriority(params.getIndexEndValue()))) {
         throw new IllegalArgumentException(
-            "When using orderByPriority(), values provided to startAt(), "
-                + "endAt(), or equalTo() must be valid priorities.");
+            "When using orderByPriority(), values provided to startAt(), startAfter(), "
+                + "endAt(), endBefore(), or equalTo() must be valid priorities.");
       }
     }
   }
@@ -115,18 +116,18 @@ private void validateQueryEndpoints(QueryParams params) {
   private void validateLimit(QueryParams params) {
     if (params.hasStart() && params.hasEnd() && params.hasLimit() && !params.hasAnchoredLimit()) {
       throw new IllegalArgumentException(
-          "Can't combine startAt(), endAt() and limit(). "
+          "Can't combine startAt(), startAfter(), endAt(), endBefore(), and limit(). "
               + "Use limitToFirst() or limitToLast() instead");
     }
   }
 
   /** This method validates that the equalTo call can be made */
   private void validateEqualToCall() {
     if (params.hasStart()) {
-      throw new IllegalArgumentException("Can't call equalTo() and startAt() combined");
+      throw new IllegalArgumentException("Cannot combine equalTo() with startAt() or startAfter()");
     }
     if (params.hasEnd()) {
-      throw new IllegalArgumentException("Can't call equalTo() and endAt() combined");
+      throw new IllegalArgumentException("Cannot combine equalTo() with endAt() or endBefore()");
     }
   }
 
@@ -289,6 +290,9 @@ public void run() {
    */
   @NonNull
   public Query startAfter(@Nullable String value) {
+    if (value != null && params.getIndex().equals(KeyIndex.getInstance())) {
+      return startAt(PushIdGenerator.successor(value));
+    }
     return startAt(value, ChildKey.getMaxName().asString());
   }
 
@@ -330,6 +334,9 @@ public Query startAfter(boolean value) {
    */
   @NonNull
   public Query startAfter(@Nullable String value, @Nullable String key) {
+    if (value != null && params.getIndex().equals(KeyIndex.getInstance())) {
+      value = PushIdGenerator.successor(value);
+    }
     Node node =
         value != null ? new StringNode(value, PriorityUtilities.NullPriority()) : EmptyNode.Empty();
     return startAfter(node, key);
@@ -488,6 +495,9 @@ private Query startAt(Node node, String key) {
    */
   @NonNull
   public Query endBefore(@Nullable String value) {
+    if (value != null && params.getIndex().equals(KeyIndex.getInstance())) {
+      return endAt(PushIdGenerator.predecessor(value));
+    }
     return endAt(value, ChildKey.getMinName().asString());
   }
 
@@ -529,6 +539,9 @@ public Query endBefore(boolean value) {
    */
   @NonNull
   public Query endBefore(@Nullable String value, @Nullable String key) {
+    if (value != null && params.getIndex().equals(KeyIndex.getInstance())) {
+      value = PushIdGenerator.predecessor(value);
+    }
     Node node =
         value != null ? new StringNode(value, PriorityUtilities.NullPriority()) : EmptyNode.Empty();
     return endBefore(node, key);

firebase-database/src/main/java/com/google/firebase/database/core/Repo.java

@@ -496,12 +496,11 @@ public Task<DataSnapshot> getValue(Query query) {
           public void run() {
             // Always check active-listener in-memory caches first. These are always at least as
             // up to date as the persistence cache.
-            Node cached =
-                serverSyncTree.calcCompleteEventCacheFromRoot(query.getPath(), new ArrayList<>());
-            if (!cached.isEmpty()) {
+            Node serverValue = serverSyncTree.getServerValue(query.getSpec());
+            if (serverValue != null) {
               source.setResult(
                   InternalHelpers.createDataSnapshot(
-                      query.getRef(), IndexedNode.from(cached, query.getSpec().getIndex())));
+                      query.getRef(), IndexedNode.from(serverValue)));
               return;
             }
             serverSyncTree.setQueryActive(query.getSpec());

firebase-database/src/main/java/com/google/firebase/database/core/SyncPoint.java

@@ -31,6 +31,7 @@
 import com.google.firebase.database.core.view.View;
 import com.google.firebase.database.core.view.ViewCache;
 import com.google.firebase.database.snapshot.ChildKey;
+import com.google.firebase.database.snapshot.EmptyNode;
 import com.google.firebase.database.snapshot.IndexedNode;
 import com.google.firebase.database.snapshot.NamedNode;
 import com.google.firebase.database.snapshot.Node;
@@ -113,11 +114,7 @@ public List<DataEvent> applyOperation(
   }
 
   /** Add an event callback for the specified query. */
-  public List<DataEvent> addEventRegistration(
-      @NotNull EventRegistration eventRegistration,
-      WriteTreeRef writesCache,
-      CacheNode serverCache) {
-    QuerySpec query = eventRegistration.getQuerySpec();
+  public View getView(QuerySpec query, WriteTreeRef writesCache, CacheNode serverCache) {
     View view = this.views.get(query.getParams());
     if (view == null) {
       // TODO: make writesCache take flag for complete server node
@@ -128,24 +125,39 @@ public List<DataEvent> addEventRegistration(
       if (eventCache != null) {
         eventCacheComplete = true;
       } else {
-        eventCache = writesCache.calcCompleteEventChildren(serverCache.getNode());
+        eventCache =
+            writesCache.calcCompleteEventChildren(
+                serverCache.getNode() != null ? serverCache.getNode() : EmptyNode.Empty());
         eventCacheComplete = false;
       }
       IndexedNode indexed = IndexedNode.from(eventCache, query.getIndex());
       ViewCache viewCache =
           new ViewCache(new CacheNode(indexed, eventCacheComplete, false), serverCache);
-      view = new View(query, viewCache);
-      // If this is a non-default query we need to tell persistence our current view of the data
-      if (!query.loadsAllData()) {
-        Set<ChildKey> allChildren = new HashSet<ChildKey>();
-        for (NamedNode node : view.getEventCache()) {
-          allChildren.add(node.getName());
-        }
-        this.persistenceManager.setTrackedQueryKeys(query, allChildren);
+      return new View(query, viewCache);
+    }
+
+    return view;
+  }
+
+  /** Add an event callback for the specified query. */
+  public List<DataEvent> addEventRegistration(
+      @NotNull EventRegistration eventRegistration,
+      WriteTreeRef writesCache,
+      CacheNode serverCache) {
+    QuerySpec query = eventRegistration.getQuerySpec();
+    View view = getView(query, writesCache, serverCache);
+    // If this is a non-default query we need to tell persistence our current view of the data
+    if (!query.loadsAllData()) {
+      Set<ChildKey> allChildren = new HashSet<ChildKey>();
+      for (NamedNode node : view.getEventCache()) {
+        allChildren.add(node.getName());
       }
+      this.persistenceManager.setTrackedQueryKeys(query, allChildren);
+    }
+    if (!this.views.containsKey(query.getParams())) {
       this.views.put(query.getParams(), view);
     }
-
+    this.views.put(query.getParams(), view);
     // This is guaranteed to exist now, we just created anything that was missing
     view.addEventRegistration(eventRegistration);
     return view.getInitialEvents(eventRegistration);

firebase-database/src/main/java/com/google/firebase/database/core/SyncTree.java

@@ -484,6 +484,60 @@ public DataSnapshot persistenceServerCache(Query query) {
         query.getRef(), persistenceManager.serverCache(query.getSpec()).getIndexedNode());
   }
 
+  @Nullable
+  public Node getServerValue(QuerySpec query) {
+    return persistenceManager.runInTransaction(
+        () -> {
+          Path path = query.getPath();
+
+          Node serverCacheNode = null;
+          boolean foundAncestorDefaultView = false;
+          // Any covering writes will necessarily be at the root, so really all we need to find is
+          // the server cache. Consider optimizing this once there's a better understanding of
+          // what actual behavior will be.
+          ImmutableTree<SyncPoint> tree = syncPointTree;
+          Path currentPath = path;
+          while (!tree.isEmpty()) {
+            SyncPoint currentSyncPoint = tree.getValue();
+            if (currentSyncPoint != null) {
+              serverCacheNode =
+                  serverCacheNode != null
+                      ? serverCacheNode
+                      : currentSyncPoint.getCompleteServerCache(currentPath);
+              foundAncestorDefaultView =
+                  foundAncestorDefaultView || currentSyncPoint.hasCompleteView();
+            }
+            ChildKey front =
+                currentPath.isEmpty() ? ChildKey.fromString("") : currentPath.getFront();
+            tree = tree.getChild(front);
+            currentPath = currentPath.popFront();
+          }
+
+          SyncPoint syncPoint = syncPointTree.get(path);
+          if (syncPoint == null) {
+            syncPoint = new SyncPoint(persistenceManager);
+            syncPointTree = syncPointTree.set(path, syncPoint);
+          } else {
+            serverCacheNode =
+                serverCacheNode != null
+                    ? serverCacheNode
+                    : syncPoint.getCompleteServerCache(Path.getEmptyPath());
+          }
+
+          CacheNode serverCache =
+              new CacheNode(
+                  IndexedNode.from(
+                      serverCacheNode != null ? serverCacheNode : EmptyNode.Empty(),
+                      query.getIndex()),
+                  serverCacheNode != null,
+                  false);
+
+          WriteTreeRef writesCache = pendingWriteTree.childWrites(path);
+          View view = syncPoint.getView(query, writesCache, serverCache);
+          return view.getCompleteNode();
+        });
+  }
+
   /** Add an event callback for the specified query. */
   public List<? extends Event> addEventRegistration(
       @NotNull final EventRegistration eventRegistration) {