2 files changed, 60 insertions, 2 deletions
diff --git a/src/main/java/com/amazon/carbonado/repo/indexed/IndexedStorage.java b/src/main/java/com/amazon/carbonado/repo/indexed/IndexedStorage.java
index 4635f2c..d115a3b 100644
--- a/src/main/java/com/amazon/carbonado/repo/indexed/IndexedStorage.java
+++ b/src/main/java/com/amazon/carbonado/repo/indexed/IndexedStorage.java
@@ -191,14 +191,26 @@ class IndexedStorage<S extends Storable> implements Storage<S>, StorageAccess<S>
         return new MergeSortBuffer<S>();
     }
 
+    public SortBuffer<S> createSortBuffer(Query.Controller controller) {
+        return new MergeSortBuffer<S>(controller);
+    }
+
     public long countAll() throws FetchException {
         return mMasterStorage.query().count();
     }
 
+    public long countAll(Query.Controller controller) throws FetchException {
+        return mMasterStorage.query().count(controller);
+    }
+
     public Cursor<S> fetchAll() throws FetchException {
         return mMasterStorage.query().fetch();
     }
 
+    public Cursor<S> fetchAll(Query.Controller controller) throws FetchException {
+        return mMasterStorage.query().fetch(controller);
+    }
+
     public Cursor<S> fetchOne(StorableIndex<S> index,
                               Object[] identityValues)
         throws FetchException
@@ -207,6 +219,15 @@ class IndexedStorage<S extends Storable> implements Storage<S>, StorageAccess<S>
         return indexInfo.fetchOne(this, identityValues);
     }
 
+    public Cursor<S> fetchOne(StorableIndex<S> index,
+                              Object[] identityValues,
+                              Query.Controller controller)
+        throws FetchException
+    {
+        ManagedIndex<S> indexInfo = (ManagedIndex<S>) mAllIndexInfoMap.get(index);
+        return indexInfo.fetchOne(this, identityValues, controller);
+    }
+
     public Query<?> indexEntryQuery(StorableIndex<S> index)
         throws FetchException
     {
@@ -221,6 +242,14 @@ class IndexedStorage<S extends Storable> implements Storage<S>, StorageAccess<S>
         return indexInfo.fetchFromIndexEntryQuery(this, indexEntryQuery);
     }
 
+    public Cursor<S> fetchFromIndexEntryQuery(StorableIndex<S> index, Query<?> indexEntryQuery,
+                                              Query.Controller controller)
+        throws FetchException
+    {
+        ManagedIndex<S> indexInfo = (ManagedIndex<S>) mAllIndexInfoMap.get(index);
+        return indexInfo.fetchFromIndexEntryQuery(this, indexEntryQuery, controller);
+    }
+
     public Cursor<S> fetchSubset(StorableIndex<S> index,
                                  Object[] identityValues,
                                  BoundaryType rangeStartBoundary,
@@ -235,6 +264,21 @@ class IndexedStorage<S extends Storable> implements Storage<S>, StorageAccess<S>
         throw new UnsupportedOperationException();
     }
 
+    public Cursor<S> fetchSubset(StorableIndex<S> index,
+                                 Object[] identityValues,
+                                 BoundaryType rangeStartBoundary,
+                                 Object rangeStartValue,
+                                 BoundaryType rangeEndBoundary,
+                                 Object rangeEndValue,
+                                 boolean reverseRange,
+                                 boolean reverseOrder,
+                                 Query.Controller controller)
+        throws FetchException
+    {
+        // This method should never be called since a query was returned by indexEntryQuery.
+        throw new UnsupportedOperationException();
+    }
+
     private void registerIndex(ManagedIndex<S> managedIndex)
         throws RepositoryException
     {
diff --git a/src/main/java/com/amazon/carbonado/repo/indexed/ManagedIndex.java b/src/main/java/com/amazon/carbonado/repo/indexed/ManagedIndex.java
index 2bb2f49..6c28619 100644
--- a/src/main/java/com/amazon/carbonado/repo/indexed/ManagedIndex.java
+++ b/src/main/java/com/amazon/carbonado/repo/indexed/ManagedIndex.java
@@ -173,6 +173,13 @@ class ManagedIndex<S extends Storable> implements IndexEntryAccessor<S> {
     Cursor<S> fetchOne(IndexedStorage storage, Object[] identityValues)
         throws FetchException
     {
+        return fetchOne(storage, identityValues, null);
+    }
+
+    Cursor<S> fetchOne(IndexedStorage storage, Object[] identityValues,
+                       Query.Controller controller)
+        throws FetchException
+    {
         Query<?> query = mSingleMatchQuery;
 
         if (query == null) {
@@ -184,13 +191,20 @@ class ManagedIndex<S extends Storable> implements IndexEntryAccessor<S> {
             mSingleMatchQuery = query = mIndexEntryStorage.query(filter);
         }
 
-        return fetchFromIndexEntryQuery(storage, query.withValues(identityValues));
+        return fetchFromIndexEntryQuery(storage, query.withValues(identityValues), controller);
     }
 
     Cursor<S> fetchFromIndexEntryQuery(IndexedStorage storage, Query<?> indexEntryQuery)
         throws FetchException
     {
-        return new IndexedCursor<S>(indexEntryQuery.fetch(), storage, mAccessor);
+        return fetchFromIndexEntryQuery(storage, indexEntryQuery, null);
+    }
+
+    Cursor<S> fetchFromIndexEntryQuery(IndexedStorage storage, Query<?> indexEntryQuery,
+                                       Query.Controller controller)
+        throws FetchException
+    {
+        return new IndexedCursor<S>(indexEntryQuery.fetch(controller), storage, mAccessor);
     }
 
     @Override