@@ -65,6 +65,32 @@ public void setIndexManager(IndexManager indexManager) {
6565 this .indexManager = indexManager ;
6666 }
6767
68+ private enum DocumentType {
69+ NO_DOCUMENT (1 ),
70+ FOUND_DOCUMENT (2 ),
71+ UNKNOWN_DOCUMENT (3 ),
72+ INVALID_DOCUMENT (4 );
73+
74+ final int dbValue ;
75+
76+ DocumentType (int dbValue ) {
77+ this .dbValue = dbValue ;
78+ }
79+
80+ static DocumentType forMutableDocument (MutableDocument document ) {
81+ if (document .isNoDocument ()) {
82+ return NO_DOCUMENT ;
83+ } else if (document .isFoundDocument ()) {
84+ return FOUND_DOCUMENT ;
85+ } else if (document .isUnknownDocument ()) {
86+ return UNKNOWN_DOCUMENT ;
87+ } else {
88+ hardAssert (!document .isValidDocument (), "MutableDocument has an unknown type" );
89+ return INVALID_DOCUMENT ;
90+ }
91+ }
92+ }
93+
6894 @ Override
6995 public void add (MutableDocument document , SnapshotVersion readTime ) {
7096 hardAssert (
@@ -77,12 +103,13 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
77103
78104 db .execute (
79105 "INSERT OR REPLACE INTO remote_documents "
80- + "(path, path_length, read_time_seconds, read_time_nanos, contents) "
81- + "VALUES (?, ?, ?, ?, ?)" ,
106+ + "(path, path_length, read_time_seconds, read_time_nanos, document_type, contents) "
107+ + "VALUES (?, ?, ?, ?, ?, ? )" ,
82108 EncodedPath .encode (documentKey .getPath ()),
83109 documentKey .getPath ().length (),
84110 timestamp .getSeconds (),
85111 timestamp .getNanoseconds (),
112+ DocumentType .forMutableDocument (document ).dbValue ,
86113 message .toByteArray ());
87114
88115 indexManager .addToCollectionParentIndex (document .getKey ().getCollectionPath ());
@@ -182,6 +209,7 @@ private Map<DocumentKey, MutableDocument> getAll(
182209 List <ResourcePath > collections ,
183210 IndexOffset offset ,
184211 int count ,
212+ @ Nullable DocumentType filterDocumentType ,
185213 @ Nullable Function <MutableDocument , Boolean > filter ,
186214 @ Nullable QueryContext context ) {
187215 Timestamp readTime = offset .getReadTime ().getTimestamp ();
@@ -192,20 +220,24 @@ private Map<DocumentKey, MutableDocument> getAll(
192220 "SELECT contents, read_time_seconds, read_time_nanos, path "
193221 + "FROM remote_documents "
194222 + "WHERE path >= ? AND path < ? AND path_length = ? "
223+ + (filterDocumentType == null ? "" : " AND document_type = ? " )
195224 + "AND (read_time_seconds > ? OR ( "
196225 + "read_time_seconds = ? AND read_time_nanos > ?) OR ( "
197226 + "read_time_seconds = ? AND read_time_nanos = ? and path > ?)) " ,
198227 collections .size (),
199228 " UNION " );
200229 sql .append ("ORDER BY read_time_seconds, read_time_nanos, path LIMIT ?" );
201230
202- Object [] bindVars = new Object [BINDS_PER_STATEMENT * collections .size () + 1 ];
231+ Object [] bindVars = new Object [BINDS_PER_STATEMENT * collections .size () + 1 + ( filterDocumentType != null ? 1 : 0 ) ];
203232 int i = 0 ;
204233 for (ResourcePath collection : collections ) {
205234 String prefixPath = EncodedPath .encode (collection );
206235 bindVars [i ++] = prefixPath ;
207236 bindVars [i ++] = EncodedPath .prefixSuccessor (prefixPath );
208237 bindVars [i ++] = collection .length () + 1 ;
238+ if (filterDocumentType != null ) {
239+ bindVars [i ++] = filterDocumentType .dbValue ;
240+ }
209241 bindVars [i ++] = readTime .getSeconds ();
210242 bindVars [i ++] = readTime .getSeconds ();
211243 bindVars [i ++] = readTime .getNanoseconds ();
@@ -235,7 +267,7 @@ private Map<DocumentKey, MutableDocument> getAll(
235267 IndexOffset offset ,
236268 int count ,
237269 @ Nullable Function <MutableDocument , Boolean > filter ) {
238- return getAll (collections , offset , count , filter , /*context*/ null );
270+ return getAll (collections , offset , count , /*filterDocumentType*/ null , filter , /*context*/ null );
239271 }
240272
241273 private void processRowInBackground (
@@ -278,6 +310,7 @@ public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
278310 Collections .singletonList (query .getPath ()),
279311 offset ,
280312 Integer .MAX_VALUE ,
313+ DocumentType .FOUND_DOCUMENT ,
281314 (MutableDocument doc ) -> query .matches (doc ) || mutatedKeys .contains (doc .getKey ()),
282315 context );
283316 }
0 commit comments