Skip to content

Commit

Permalink
feat: support generating shell syntax for $group
Browse files Browse the repository at this point in the history
  • Loading branch information
himanshusinghs committed Jan 30, 2025
1 parent facd7e2 commit 7c23a54
Show file tree
Hide file tree
Showing 11 changed files with 1,797 additions and 58 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,9 @@
import com.mongodb.client.model.Sorts;
import org.bson.Document;
import org.bson.types.ObjectId;
import org.springframework.data.mongodb.core.aggregation.Fields;

import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.UUID;
import java.util.*;

public class JavaDriverRepository {
private static final String IMDB_VOTES = "imdb.votes";
Expand Down Expand Up @@ -94,28 +92,14 @@ public List<Document> queryMoviesByYear(int year) {
.getCollection("movies")
.aggregate(
List.of(
Aggregates.match(
Filters.eq("year", year)
),
Aggregates.group(
"newField",
Accumulators.avg("test", "$year"),
Accumulators.sum("test2", "$year"),
Accumulators.bottom("field", Sorts.ascending("year"), "$year")
),
Aggregates.project(
Projections.fields(
Projections.include("year", "plot")
null,
Accumulators.topN(
"3fields",
Sorts.descending("year"),
"fields",
3
)
),
Aggregates.sort(
Sorts.orderBy(
Sorts.ascending("asd", "qwe")
)
),
Aggregates.unwind(
"awards.wins",
new UnwindOptions()
)
)
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,17 +35,7 @@ private List<Movie> allMoviesWithRatingAtLeast(int rating) {
private List<Movie> allMoviesWithRatingAtLeastAgg(int rating) {
return template.aggregate(
Aggregation.newAggregation(
Aggregation.match(where( "tomatoes.viewer.rating").gte(rating)),
Aggregation.project("fieldA").andInclude("fieldB").andExclude("fieldC"),
Aggregation.sort(Sort.Direction.ASC, "asd"),
Aggregation.sort(Sort.by("rated", "qwe")),
Aggregation.sort(Sort.by(Sort.Direction.ASC, "rates")),
Aggregation.sort(Sort.by(Sort.Order.by("rated"), Sort.Order.by("ratedd"))),
Aggregation.sort(Sort.by(List.of(Sort.Order.by("rated"), Sort.Order.by("rateds")))),
Aggregation.addFields().addFieldWithValueOf("addedField", "value").build(),
Aggregation.addFields().addFieldWithValueOf("addedField", Fields.field("qwe")).build(),
Aggregation.addFields().addField("addedField").withValueOf("rateds").build(),
Aggregation.addFields().addField("addedField").withValueOf(Fields.field("asd")).build()
Aggregation.group()
),
Movie.class,
Movie.class
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ fun <S> MongoshBackend.emitAggregateBody(node: Node<S>, queryContext: QueryConte
Name.ADD_FIELDS -> emitAddFieldsStage(stage)
Name.UNWIND -> emitUnwindStage(stage)
Name.SORT -> emitSortStage(stage)
Name.GROUP -> emitGroupStage(stage)
else -> {}
}
emitObjectValueEnd(long = true)
Expand All @@ -49,7 +50,12 @@ internal fun <S> MongoshBackend.emitAsFieldValueDocument(nodes: List<Node<S>>, i
val field = node.component<HasFieldReference<S>>() ?: continue
val value = node.component<HasValueReference<S>>() ?: continue

emitObjectKey(resolveFieldReference(field))
emitObjectKey(
resolveFieldReference(
fieldRef = field,
fieldUsedAsValue = false,
)
)
emitContextValue(resolveValueReference(value, field))
emitObjectValueEnd(long = isLong)
}
Expand All @@ -63,6 +69,7 @@ private val NON_DESTRUCTIVE_STAGES = setOf(
Name.ADD_FIELDS,
Name.UNWIND,
Name.SORT,
Name.GROUP,
)

private fun <S> Node<S>.isNotDestructive(): Boolean {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,176 @@
package com.mongodb.jbplugin.dialects.mongosh.aggr

import com.mongodb.jbplugin.dialects.mongosh.backend.MongoshBackend
import com.mongodb.jbplugin.dialects.mongosh.query.resolveFieldReference
import com.mongodb.jbplugin.dialects.mongosh.query.resolveValueReference
import com.mongodb.jbplugin.mql.Node
import com.mongodb.jbplugin.mql.components.HasAccumulatedFields
import com.mongodb.jbplugin.mql.components.HasFieldReference
import com.mongodb.jbplugin.mql.components.HasLimit
import com.mongodb.jbplugin.mql.components.HasSorts
import com.mongodb.jbplugin.mql.components.HasValueReference
import com.mongodb.jbplugin.mql.components.Name
import com.mongodb.jbplugin.mql.components.Named

internal fun <S>MongoshBackend.emitGroupStage(node: Node<S>): MongoshBackend {
val idFieldReference = node.component<HasFieldReference<S>>()
val idValueReference = node.component<HasValueReference<S>>()

val accumulatedFields = node.component<HasAccumulatedFields<S>>()?.children ?: emptyList()
val emitLongQuery = idValueReference.isLongIdValueReference() || accumulatedFields.size > 2

// "{"
emitObjectStart(long = emitLongQuery)
// "{ $group : "
emitObjectKey(registerConstant('$' + "group"))
// "{ $group : { "
emitObjectStart(long = emitLongQuery)
if (idFieldReference != null && idValueReference != null) {
// "{ $group : { _id : null, "
emitAsFieldValueDocument(listOf(node), isLong = emitLongQuery)
// "{ $group : { _id : null, totalCount: { $sum : 1 }, "
emitAccumulatedFields(accumulatedFields, emitLongQuery)
}
// "{ $group : { _id : null, totalCount: { $sum : 1 }, }"
emitObjectEnd(long = emitLongQuery)
// "{ $group : { _id : null, totalCount: { $sum : 1 }, } }"
emitObjectEnd(long = emitLongQuery)

return this
}

private fun <S>HasValueReference<S>?.isLongIdValueReference(): Boolean {
val computedReference = this?.reference as? HasValueReference.Computed<S>
?: return false
val fieldReferences = computedReference.type.expression.components<HasFieldReference<S>>()
return fieldReferences.size >= 3
}

private fun <S>MongoshBackend.emitAccumulatedFields(
accumulatedFields: List<Node<S>>,
emitLongQuery: Boolean
): MongoshBackend {
for (accumulatedField in accumulatedFields) {
val accumulator = accumulatedField.component<Named>() ?: continue
when (accumulator.name) {
Name.SUM,
Name.AVG,
Name.MIN,
Name.MAX,
Name.FIRST,
Name.LAST,
Name.PUSH,
Name.ADD_TO_SET -> {
emitKeyValueAccumulator(
accumulator,
accumulatedField,
emitLongQuery
)
// "{ $group : { _id : null, totalCount : { $sum : 1 }
emitObjectValueEnd(long = emitLongQuery)
}
Name.TOP,
Name.TOP_N,
Name.BOTTOM,
Name.BOTTOM_N -> {
emitTopBottomAccumulator(
accumulator,
accumulatedField,
emitLongQuery
)
// "{ $group : { _id : null, totalCount : { $top : { sortBy: { year: -1 }, "title" } }
emitObjectValueEnd(long = emitLongQuery)
}
else -> continue
}
}
return this
}

/**
* Emits for the following accumulators
* sum, avg, first, last, max, min, push, addToSet
*/
private fun <S>MongoshBackend.emitKeyValueAccumulator(
accumulator: Named,
accumulatedField: Node<S>,
emitLongQuery: Boolean,
): MongoshBackend {
val fieldRef = accumulatedField.component<HasFieldReference<S>>() ?: return this
val valueRef = accumulatedField.component<HasValueReference<S>>() ?: return this

// "{ $group : { _id : null, totalCount :
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
// "{ $group : { _id : null, totalCount : {
emitObjectStart(long = emitLongQuery)
// "{ $group : { _id : null, totalCount : { $sum :
emitObjectKey(registerConstant('$' + accumulator.name.canonical))
// "{ $group : { _id : null, totalCount : { $sum : 1
emitContextValue(resolveValueReference(valueRef, fieldRef))
// "{ $group : { _id : null, totalCount : { $sum : 1 }
emitObjectEnd(long = emitLongQuery)
return this
}

private fun <S>MongoshBackend.emitTopBottomAccumulator(
accumulator: Named,
accumulatedField: Node<S>,
emitLongQuery: Boolean,
): MongoshBackend {
val fieldRef = accumulatedField.component<HasFieldReference<S>>() ?: return this
val valueRef = accumulatedField.component<HasValueReference<S>>() ?: return this
val sorts = accumulatedField.component<HasSorts<S>>()?.children ?: emptyList()
val limit = accumulatedField.component<HasLimit>()?.limit
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
// "{"
emitObjectStart(long = emitLongQuery)
// "{ $top : "
emitObjectKey(registerConstant('$' + accumulator.name.canonical))
// "{ $top : {"
emitObjectStart(long = emitLongQuery)
// "{ $top : { "sortBy" : "
emitObjectKey(registerConstant("sortBy"))
// "{ $top : { "sortBy" : { "
emitObjectStart(long = emitLongQuery)
// "{ $top : { "sortBy" : { "field" : 1,"
emitAsFieldValueDocument(sorts, emitLongQuery)
// "{ $top : { "sortBy" : { "field" : 1, }"
emitObjectEnd(long = emitLongQuery)
// "{ $top : { "sortBy" : { "field" : 1, }, "
emitObjectValueEnd(long = emitLongQuery)

// "{ $top : { "sortBy" : { "field" : 1, }, output : "
emitObjectKey(registerConstant("output"))
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField""
emitContextValue(resolveValueReference(valueRef, fieldRef))
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", "
emitObjectValueEnd(long = emitLongQuery)

if (limit != null) {
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", n : "
emitObjectKey(registerConstant("n"))
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", n : 3"
emitContextValue(registerConstant(limit))
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", n : 3, "
emitObjectValueEnd(long = emitLongQuery)
}

// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", } " or
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", n : 3, }"
emitObjectEnd(long = emitLongQuery)

// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", } }" or
// "{ $top : { "sortBy" : { "field" : 1, }, output : "$someField", n : 3, } }"
emitObjectEnd(long = emitLongQuery)
return this
}
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,12 @@ internal fun <S> MongoshBackend.emitUnwindStage(node: Node<S>): MongoshBackend {

emitObjectStart()
emitObjectKey(registerConstant('$' + "unwind"))
emitContextValue(resolveFieldReference(unwindField))
emitContextValue(
resolveFieldReference(
fieldRef = unwindField,
fieldUsedAsValue = true,
)
)
emitObjectEnd()

return this
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ class DefaultContext : Context {
val existingVariable = variables.getOrElse(cleanName) { null }

if (existingVariable != null && existingVariable.value == null) {
// already exists, generate a new name
// already exists, generate a new name
val nameWithoutCounter = if (cleanName.matches(endsWithNumber)) {
cleanName.replace(endsWithNumber, "")
} else {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,12 @@ fun <S> MongoshBackend.emitQueryFilter(node: Node<S>, firstCall: Boolean = false
if (firstCall) {
emitObjectStart(long = isLong)
}
emitObjectKey(resolveFieldReference(fieldRef))
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
emitContextValue(resolveValueReference(valueRef, fieldRef))
if (firstCall) {
emitObjectEnd(long = isLong)
Expand All @@ -56,7 +61,12 @@ fun <S> MongoshBackend.emitQueryFilter(node: Node<S>, firstCall: Boolean = false
emitObjectStart(long = isLong)
}
if (fieldRef != null) {
emitObjectKey(resolveFieldReference(fieldRef))
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
}

if (valueRef != null) {
Expand Down Expand Up @@ -85,7 +95,12 @@ fun <S> MongoshBackend.emitQueryFilter(node: Node<S>, firstCall: Boolean = false
}

if (fieldRef != null) {
emitObjectKey(resolveFieldReference(fieldRef))
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
}

emitObjectStart()
Expand Down Expand Up @@ -154,7 +169,12 @@ fun <S> MongoshBackend.emitQueryFilter(node: Node<S>, firstCall: Boolean = false
}

// emit field name first
emitObjectKey(resolveFieldReference(fieldRef))
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
// emit the $not
emitObjectStart()
emitObjectKey(registerConstant('$' + "not"))
Expand All @@ -176,7 +196,12 @@ fun <S> MongoshBackend.emitQueryFilter(node: Node<S>, firstCall: Boolean = false
if (firstCall) {
emitObjectStart(long = isLong)
}
emitObjectKey(resolveFieldReference(fieldRef))
emitObjectKey(
resolveFieldReference(
fieldRef = fieldRef,
fieldUsedAsValue = false,
)
)
emitObjectStart(long = isLong)
emitObjectKey(registerConstant('$' + named.name.canonical))
emitContextValue(resolveValueReference(valueRef, fieldRef))
Expand Down
Loading

0 comments on commit 7c23a54

Please sign in to comment.