/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.jackrabbit.oak.plugins.document.mongo;

import java.io.Closeable;
import java.util.Queue;

import com.google.common.collect.AbstractIterator;
import com.google.common.collect.Queues;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.QueryBuilder;
import org.apache.jackrabbit.oak.plugins.blob.ReferencedBlob;
import org.apache.jackrabbit.oak.plugins.document.BlobCollector;
import org.apache.jackrabbit.oak.plugins.document.Collection;
import org.apache.jackrabbit.oak.plugins.document.DocumentNodeStore;
import org.apache.jackrabbit.oak.plugins.document.NodeDocument;

public class MongoBlobReferenceIterator extends AbstractIterator<ReferencedBlob> implements Closeable {
    private final MongoDocumentStore documentStore;
    private final BlobCollector blobCollector;
    private final Queue<ReferencedBlob> blobs = Queues.newArrayDeque();

    private DBCursor cursor;

    public MongoBlobReferenceIterator(DocumentNodeStore nodeStore,
                                      MongoDocumentStore documentStore) {
        this.documentStore = documentStore;
        this.blobCollector = new BlobCollector(nodeStore);
    }

    @Override
    protected ReferencedBlob computeNext() {
        if (blobs.isEmpty()) {
            loadBatch();
        }
        if (!blobs.isEmpty()) {
            return blobs.remove();
        }
        return endOfData();
    }

    private void loadBatch() {
        initializeCursor();
        //Some node which have the '_bin' flag set might not have any binaries in it
        //so move forward if blobs is still empty and cursor has more elements
        while (cursor.hasNext() && blobs.isEmpty()) {
            collectBinaries(documentStore.convertFromDBObject(Collection.NODES, cursor.next()));
        }
    }

    private void collectBinaries(NodeDocument nodeDocument) {
        blobCollector.collect(nodeDocument, blobs);
    }

    private void initializeCursor() {
        if (cursor == null) {
            DBObject query = QueryBuilder.start(NodeDocument.HAS_BINARY_FLAG)
                    .is(NodeDocument.HAS_BINARY_VAL)
                    .get();
            //TODO It currently prefers secondary. Would that be Ok?
            cursor = getNodeCollection().find(query)
                    .setReadPreference(documentStore.getConfiguredReadPreference(Collection.NODES));
        }
    }

    private DBCollection getNodeCollection() {
        return documentStore.getDBCollection(Collection.NODES);
    }

    @Override
    public void close() {
        if (cursor != null) {
            cursor.close();
        }
    }
}

@@ -872,268 +872,91 @@
 ort 
-java.io.Closeable;%0Aimport java.util.Queue;%0A%0Aimport com.google.common.collect.AbstractIterator;%0Aimport com.google.common.collect.Queues;%0Aimport com.mongodb.DBCollection;%0Aimport com.mongodb.DBCursor;%0Aimport com.mongodb.DBObject;%0Aimport com.mongodb.QueryBuilde
+static com.google.common.collect.Iterables.transform;%0A%0Aimport java.util.Iterato
 r;%0A
+%0A
 impo
@@ -996,13 +996,21 @@
 ins.
-b
+document.B
 lob
-.
 Refe
@@ -1018,75 +1018,13 @@
 ence
-dBlob;%0Aimport org.apache.jackrabbit.oak.plugins.document.BlobCollec
+Itera
 tor;
@@ -1223,113 +1223,288 @@
 nt;%0A
-%0Apublic class MongoBlobReferenceIterator extends AbstractIterator%3CReferencedBlob%3E implements Closeable
+import org.apache.jackrabbit.oak.plugins.document.util.CloseableIterable;%0A%0Aimport com.google.common.base.Function;%0Aimport com.mongodb.DBCursor;%0Aimport com.mongodb.DBObject;%0Aimport com.mongodb.QueryBuilder;%0A%0Apublic class MongoBlobReferenceIterator extends BlobReferenceIterator
  %7B%0A
+%0A
     
@@ -1554,157 +1554,8 @@
 ore;
-%0A    private final BlobCollector blobCollector;%0A    private final Queue%3CReferencedBlob%3E blobs = Queues.newArrayDeque();%0A%0A    private DBCursor cursor;
 %0A%0A  
@@ -1622,46 +1622,8 @@
 ore,
-%0A                                     
  Mon
@@ -1667,42 +1667,24 @@
     
-this.documentStore = document
+super(node
 Store
+)
 ;%0A  
@@ -1698,52 +1698,37 @@
 his.
-blobCollec
+documentS
 tor
+e
  = 
-new BlobCollector(node
+document
 Store
-)
 ;%0A  
@@ -1755,810 +1755,73 @@
    p
-rotected ReferencedBlob computeNext() %7B%0A        if (blobs.isEmpty()) %7B%0A            loadBatch();%0A        %7D%0A        if (!blobs.isEmpty()) %7B%0A            return blobs.remove();%0A        %7D%0A        return endOfData();%0A    %7D%0A%0A    private void loadBatch() %7B%0A        initializeCursor();%0A        //Some node which have the '_bin' flag set might not have any binaries in it%0A        //so move forward if blobs is still empty and cursor has more elements%0A        while (cursor.hasNext() && blobs.isEmpty()) %7B%0A            collectBinaries(documentStore.convertFromDBObject(Collection.NODES, cursor.next()));%0A        %7D%0A    %7D%0A%0A    private void collectBinaries(NodeDocument nodeDocument) %7B%0A        blobCollector.collect(nodeDocument, blobs);%0A    %7D%0A%0A    private void initializeCursor() %7B%0A        if (cursor == null) %7B%0A    
+ublic Iterator%3CNodeDocument%3E getIteratorOverDocsWithBinaries() %7B%0A
     
@@ -1889,37 +1889,16 @@
 RY_FLAG)
-%0A                    
 .is(Node
@@ -1921,37 +1921,16 @@
 ARY_VAL)
-%0A                    
 .get();%0A
@@ -1941,14 +1941,11 @@
     
-    
 //
+ 
 TODO
@@ -1999,27 +1999,32 @@
 Ok?%0A        
-   
+DBCursor
  cursor = ge
@@ -2021,23 +2021,46 @@
 ursor = 
-getNode
+documentStore.getDBCollection(
 Collecti
@@ -2061,17 +2061,22 @@
 llection
-(
+.NODES
 ).find(q
@@ -2081,20 +2081,16 @@
 (query)%0A
-    
         
@@ -2182,80 +2182,197 @@
 ));%0A
-        %7D
 %0A    
-%7D%0A%0A
     
-private DBCollection getNodeCollection() %7B%0A
+return CloseableIterable.wrap(transform(cursor, new Function%3CDBObject, NodeDocument%3E() %7B%0A            @Override%0A            public NodeDocument apply(DBObject input) %7B%0A     
         
 retu
@@ -2363,24 +2363,27 @@
             
+   
 return docum
@@ -2395,119 +2395,67 @@
 ore.
-getDBCollection(Collection.NODES);%0A    %7D%0A%0A    @Override%0A    public void close() %7B%0A        if (cursor != null) %7B
+convertFromDBObject(Collection.NODES, input);%0A            %7D
 %0A   
@@ -2463,37 +2463,32 @@
     
-   
+%7D),
  cursor
-.close();%0A        %7D
+).iterator();%0A
 %0A   

Refactor MongoBlobReferenceIterator