Project

General

Profile

1
package eu.dnetlib.oai.mongo;
2

    
3
import java.util.Iterator;
4
import java.util.function.Function;
5

    
6
import com.google.common.collect.Lists;
7
import com.mongodb.DBObject;
8
import com.mongodb.client.MongoCursor;
9

    
10
import eu.dnetlib.oai.Cursor;
11
import eu.dnetlib.oai.info.RecordInfo;
12

    
13
public class DNetOAIMongoCursor implements Cursor {
14

    
15
	/**
16
	 * Underlying mongo cursor.
17
	 */
18
	private MongoCursor<DBObject> dbCursor;
19
	private int size = 0;
20
	/**
21
	 * Function to apply to records before delivering.
22
	 */
23
	private Function<String, String> function;
24

    
25
	/**
26
	 * true if the RecordInfo returned by this Cursor must include the record body, false otherwise.
27
	 */
28
	private boolean bodyIncluded;
29

    
30
	private RecordInfoGenerator recordInfoGenerator;
31
	private MetadataExtractor metadataExtractor;
32
	private ProvenanceExtractor provenanceExtractor;
33

    
34
	public DNetOAIMongoCursor() {
35
		super();
36
	}
37

    
38
	public DNetOAIMongoCursor(final MongoCursor<DBObject> dbCursor, final boolean bodyIncluded, final RecordInfoGenerator recordInfoGenerator,
39
			final MetadataExtractor metadataExtractor) {
40
		this(dbCursor, null, bodyIncluded, recordInfoGenerator, metadataExtractor);
41
	}
42

    
43
	public DNetOAIMongoCursor(final MongoCursor<DBObject> dbCursor, final Function<String, String> function, final boolean bodyIncluded,
44
			final RecordInfoGenerator recordInfoGenerator, final MetadataExtractor metadataExtractor) {
45
		super();
46
		this.dbCursor = dbCursor;
47
		this.size = 0;
48
		this.function = function;
49
		this.bodyIncluded = bodyIncluded;
50
		this.recordInfoGenerator = recordInfoGenerator;
51
		this.metadataExtractor = metadataExtractor;
52
	}
53

    
54
	/**
55
	 * {@inheritDoc}
56
	 */
57
	@Override
58
	public int count() {
59
		// I can do it because MongoCursor are always created from queries with "limit", so I do not expect the creation of the list to
60
		// explode
61
		// to not exagerate, I'll get the size only if the current size is 0
62
		if (this.size == 0) {
63
			this.size = Lists.newArrayList(this.dbCursor).size();
64
		}
65
		return this.size;
66
	}
67

    
68
	/**
69
	 * {@inheritDoc}
70
	 *
71
	 * @see Iterable#iterator()
72
	 */
73
	@Override
74
	public Iterator<RecordInfo> iterator() {
75

    
76
		return new Iterator<RecordInfo>() {
77

    
78
			@Override
79
			public boolean hasNext() {
80
				return DNetOAIMongoCursor.this.dbCursor.hasNext();
81
			}
82

    
83
			@Override
84
			public RecordInfo next() {
85
				final DBObject res = DNetOAIMongoCursor.this.dbCursor.next();
86
				final RecordInfo info = DNetOAIMongoCursor.this.recordInfoGenerator.transformDBObject(res, DNetOAIMongoCursor.this.bodyIncluded);
87
				if ((DNetOAIMongoCursor.this.function != null) && DNetOAIMongoCursor.this.bodyIncluded && (info != null)) {
88
					info.setMetadata(DNetOAIMongoCursor.this.function.apply(info.getMetadata()));
89
				}
90
				return info;
91
			}
92

    
93
			@Override
94
			public void remove() {
95
				throw new UnsupportedOperationException();
96
			}
97

    
98
		};
99
	}
100

    
101
	public Function<String, String> getFunction() {
102
		return this.function;
103
	}
104

    
105
	public void setFunction(final Function<String, String> function) {
106
		this.function = function;
107
	}
108

    
109
	public MongoCursor<DBObject> getDbCursor() {
110
		return this.dbCursor;
111
	}
112

    
113
	public void setDbCursor(final MongoCursor<DBObject> dbCursor) {
114
		this.dbCursor = dbCursor;
115
	}
116

    
117
	@Override
118
	public boolean isBodyIncluded() {
119
		return this.bodyIncluded;
120
	}
121

    
122
	@Override
123
	public void setBodyIncluded(final boolean bodyIncluded) {
124
		this.bodyIncluded = bodyIncluded;
125
	}
126

    
127
	public RecordInfoGenerator getRecordInfoGenerator() {
128
		return this.recordInfoGenerator;
129
	}
130

    
131
	public void setRecordInfoGenerator(final RecordInfoGenerator recordInfoGenerator) {
132
		this.recordInfoGenerator = recordInfoGenerator;
133
	}
134

    
135
	public MetadataExtractor getMetadataExtractor() {
136
		return this.metadataExtractor;
137
	}
138

    
139
	public void setMetadataExtractor(final MetadataExtractor metadataExtractor) {
140
		this.metadataExtractor = metadataExtractor;
141
	}
142

    
143
	public ProvenanceExtractor getProvenanceExtractor() {
144
		return this.provenanceExtractor;
145
	}
146

    
147
	public void setProvenanceExtractor(final ProvenanceExtractor provenanceExtractor) {
148
		this.provenanceExtractor = provenanceExtractor;
149
	}
150

    
151
}
(1-1/6)