Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support generating shell syntax for $group INTELLIJ-197 #134

Merged
merged 4 commits into from
Jan 30, 2025
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA.
## [Unreleased]

### Added
* [INTELLIJ-197](https://jira.mongodb.org/browse/INTELLIJ-197) Add support for generating shell syntax for $group stage and supported accumulators when running queries.
* [INTELLIJ-198](https://jira.mongodb.org/browse/INTELLIJ-198) New modal to provide default values when generating queries with unknown runtime expressions.
* [INTELLIJ-175](https://jira.mongodb.org/browse/INTELLIJ-175) Add support for parsing, inspecting and autocompleting in a group stage written using `Aggregation.group` and chained `GroupOperation`s using `sum`, `avg`, `first`, `last`, `max`, `min`, `push` and `addToSet`.
* [INTELLIJ-196](https://jira.mongodb.org/browse/INTELLIJ-196) Add support for $sort when generating the query into DataGrip.
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,
)
)
Comment on lines +53 to +58
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kmruiz I adopted this approach to make MongoshBackend understand when to prefix the field name with a $ and when not to. The idea is if a field reference is used in the value place then it is supposed to be true and the emitted field will be prefixed with $ otherwise not.

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,
himanshusinghs marked this conversation as resolved.
Show resolved Hide resolved
)

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
Loading