Skip to content

Commit 8802a4b

Browse files
committed
Ability to specify union
JAVA-3472
1 parent 7381f90 commit 8802a4b

File tree

4 files changed

+135
-0
lines changed

4 files changed

+135
-0
lines changed

driver-core/src/main/com/mongodb/client/model/Aggregates.java

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,20 @@ public static <TExpression> Bson group(@Nullable final TExpression id, final Lis
377377
return new GroupStage<TExpression>(id, fieldAccumulators);
378378
}
379379

380+
/**
381+
* Creates a $unionWith pipeline stage.
382+
*
383+
* @param collection the name of the collection in the same database to perform the union with.
384+
* @param pipeline the pipeline to run on the union.
385+
* @return the $unionWith pipeline stage
386+
* @mongodb.driver.manual reference/operator/aggregation/unionWith/ $unionWith
387+
* @mongodb.server.release 4.4
388+
* @since 4.1
389+
*/
390+
public static Bson unionWith(final String collection, final List<? extends Bson> pipeline) {
391+
return new UnionWithStage(collection, pipeline);
392+
}
393+
380394
/**
381395
* Creates a $unwind pipeline stage for the specified field name, which must be prefixed by a {@code '$'} sign.
382396
*
@@ -1376,6 +1390,70 @@ public String toString() {
13761390
}
13771391
}
13781392

1393+
private static final class UnionWithStage implements Bson {
1394+
private final String collection;
1395+
private final List<? extends Bson> pipeline;
1396+
1397+
private UnionWithStage(final String collection, final List<? extends Bson> pipeline) {
1398+
this.collection = collection;
1399+
this.pipeline = pipeline;
1400+
}
1401+
1402+
@Override
1403+
public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> tDocumentClass, final CodecRegistry codecRegistry) {
1404+
BsonDocumentWriter writer = new BsonDocumentWriter(new BsonDocument());
1405+
1406+
writer.writeStartDocument();
1407+
1408+
writer.writeStartDocument("$unionWith");
1409+
writer.writeString("coll", collection);
1410+
1411+
writer.writeName("pipeline");
1412+
writer.writeStartArray();
1413+
for (Bson stage : pipeline) {
1414+
BuildersHelper.encodeValue(writer, stage, codecRegistry);
1415+
}
1416+
writer.writeEndArray();
1417+
1418+
writer.writeEndDocument();
1419+
1420+
return writer.getDocument();
1421+
}
1422+
1423+
@Override
1424+
public boolean equals(final Object o) {
1425+
if (this == o) {
1426+
return true;
1427+
}
1428+
if (o == null || getClass() != o.getClass()) {
1429+
return false;
1430+
}
1431+
1432+
UnionWithStage that = (UnionWithStage) o;
1433+
1434+
if (!collection.equals(that.collection)) {
1435+
return false;
1436+
}
1437+
return pipeline != null ? !pipeline.equals(that.pipeline) : that.pipeline != null;
1438+
}
1439+
1440+
@Override
1441+
public int hashCode() {
1442+
int result = collection.hashCode();
1443+
result = 31 * result + (pipeline != null ? pipeline.hashCode() : 0);
1444+
return result;
1445+
}
1446+
1447+
@Override
1448+
public String toString() {
1449+
return "Stage{"
1450+
+ "name='$unionWith'"
1451+
+ ", collection='" + collection + '\''
1452+
+ ", pipeline=" + pipeline
1453+
+ '}';
1454+
}
1455+
}
1456+
13791457
private Aggregates() {
13801458
}
13811459
}

driver-core/src/test/functional/com/mongodb/client/model/AggregatesFunctionalSpecification.groovy

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ import static com.mongodb.client.model.Aggregates.sample
5656
import static com.mongodb.client.model.Aggregates.skip
5757
import static com.mongodb.client.model.Aggregates.sort
5858
import static com.mongodb.client.model.Aggregates.sortByCount
59+
import static com.mongodb.client.model.Aggregates.unionWith
5960
import static com.mongodb.client.model.Aggregates.unwind
6061
import static com.mongodb.client.model.Filters.eq
6162
import static com.mongodb.client.model.Filters.exists
@@ -981,4 +982,41 @@ class AggregatesFunctionalSpecification extends OperationFunctionalSpecification
981982
then:
982983
results == [Document.parse('{b: 1, _id: 7}')]
983984
}
985+
986+
@IgnoreIf({ !serverVersionAtLeast(4, 3) })
987+
def '$unionWith'() {
988+
given:
989+
def coll1Helper = getCollectionHelper(new MongoNamespace(getDatabaseName(), 'coll1'))
990+
def coll2Helper = getCollectionHelper(new MongoNamespace(getDatabaseName(), 'coll2'))
991+
992+
coll1Helper.drop()
993+
coll2Helper.drop()
994+
995+
coll1Helper.insertDocuments(
996+
Document.parse('{ "name1" : "almonds" }'),
997+
Document.parse('{ "name1" : "cookies" }'))
998+
999+
coll2Helper.insertDocuments(
1000+
Document.parse('{ "name2" : "cookies" }'),
1001+
Document.parse('{ "name2" : "cookies" }'),
1002+
Document.parse('{ "name2" : "pecans" }'))
1003+
1004+
def pipeline = asList(match(eq('name2', 'cookies')), project(fields(excludeId(), computed('name', '$name2'))),
1005+
sort(ascending('name')))
1006+
1007+
when:
1008+
def results = coll1Helper.aggregate([project(fields(excludeId(), computed('name', '$name1'))),
1009+
unionWith('coll2', pipeline)])
1010+
1011+
then:
1012+
results == [
1013+
Document.parse('{ name: "almonds" }'),
1014+
Document.parse('{ name: "cookies" }'),
1015+
Document.parse('{ name: "cookies" }'),
1016+
Document.parse('{ name: "cookies" }') ]
1017+
1018+
cleanup:
1019+
coll1Helper?.drop()
1020+
coll2Helper?.drop()
1021+
}
9841022
}

driver-core/src/test/unit/com/mongodb/client/model/AggregatesSpecification.groovy

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ import static com.mongodb.client.model.Aggregates.sample
5656
import static com.mongodb.client.model.Aggregates.skip
5757
import static com.mongodb.client.model.Aggregates.sort
5858
import static com.mongodb.client.model.Aggregates.sortByCount
59+
import static com.mongodb.client.model.Aggregates.unionWith
5960
import static com.mongodb.client.model.Aggregates.unwind
6061
import static com.mongodb.client.model.BsonHelper.toBson
6162
import static com.mongodb.client.model.Filters.eq
@@ -327,6 +328,13 @@ class AggregatesSpecification extends Specification {
327328
toBson(skip(5)) == parse('{ $skip : 5 }')
328329
}
329330

331+
def 'should render $unionWith'() {
332+
expect:
333+
List<Bson> pipeline = asList(match(expr(new Document('$eq', asList('x', '1')))))
334+
toBson(unionWith('with', pipeline)) ==
335+
parse('''{ $unionWith : { coll: "with", pipeline : [{ $match : { $expr: { $eq : [ "x" , "1" ]}}}] }}''')
336+
}
337+
330338
def 'should render $unwind'() {
331339
expect:
332340
toBson(unwind('$sizes')) == parse('{ $unwind : "$sizes" }')

driver-scala/src/main/scala/org/mongodb/scala/model/Aggregates.scala

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,4 +431,15 @@ object Aggregates {
431431
def merge(namespace: MongoNamespace, mergeOptions: MergeOptions): Bson =
432432
JAggregates.merge(namespace, mergeOptions.wrapped)
433433

434+
/**
435+
* Creates a `\$unionWith` pipeline stage.
436+
*
437+
* @param collection the name of the collection in the same database to perform the union with.
438+
* @param pipeline the pipeline to run on the union.
439+
* @return the \$unionWith` pipeline stage
440+
* @see [[http://docs.mongodb.org/manual/reference/operator/aggregation/unionWith/]]
441+
*/
442+
def unionWith(collection: String, pipeline: Bson*): Bson =
443+
JAggregates.unionWith(collection, pipeline.asJava)
444+
434445
}

0 commit comments

Comments
 (0)