Skip to content

Commit

Permalink
chore: project!
Browse files Browse the repository at this point in the history
  • Loading branch information
kmruiz committed Jan 22, 2025
1 parent 970a369 commit c40b321
Show file tree
Hide file tree
Showing 6 changed files with 187 additions and 31 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,6 @@ import com.mongodb.jbplugin.meta.service
import com.mongodb.jbplugin.mql.Node
import com.mongodb.jbplugin.mql.QueryContext
import com.mongodb.jbplugin.mql.components.HasSourceDialect
import com.mongodb.jbplugin.mql.components.IsCommand
import com.mongodb.jbplugin.mql.parser.components.whenHasAnyCommand
import com.mongodb.jbplugin.mql.parser.components.whenIsCommand
import com.mongodb.jbplugin.mql.parser.first
import com.mongodb.jbplugin.mql.parser.map
import com.mongodb.jbplugin.mql.parser.parse
import com.mongodb.jbplugin.observability.TelemetryEvent
import com.mongodb.jbplugin.observability.TelemetryEvent.QueryRunEvent.Console
import com.mongodb.jbplugin.observability.probe.QueryRunProbe
Expand All @@ -69,10 +63,6 @@ internal object RunQueryCodeAction : MongoDbCodeAction {
query: Node<PsiElement>,
formatter: DialectFormatter
): LineMarkerInfo<PsiElement>? {
if (!shouldShowRunGutterIcon(query)) {
return null
}

return LineMarkerInfo(
query.sourceForMarker,
query.sourceForMarker.textRange,
Expand Down Expand Up @@ -124,14 +114,6 @@ internal object RunQueryCodeAction : MongoDbCodeAction {
)
}

private fun shouldShowRunGutterIcon(node: Node<PsiElement>): Boolean {
return first(
whenIsCommand<PsiElement>(IsCommand.CommandType.AGGREGATE).map { false },
whenHasAnyCommand<PsiElement>().map { true }
).parse(node) // if we couldn't parse the query, don't show the gutter icon
.orElse { false }
}

private fun openDataGripConsole(
query: Node<PsiElement>,
newDataSource: LocalDataSource,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import com.mongodb.client.model.Projections;
import com.mongodb.client.model.Sorts;
import org.bson.Document;
import org.bson.types.ObjectId;

import java.util.ArrayList;
import java.util.List;
Expand Down Expand Up @@ -57,13 +58,13 @@ public List<Document> findMoviesByYear(int year) {
.into(new ArrayList<>());
}

public Document queryMovieById(String id) {
public Document queryMovieById(ObjectId id) {
return client
.getDatabase("sample_mflix")
.getCollection("movies")
.aggregate(List.of(Aggregates.match(
Filters.eq(id)
), Aggregates.unwind("year")))
), Aggregates.project(Projections.include("title", "year"))))
.first();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ fun <S> MongoshBackend.emitAggregateBody(node: Node<S>, queryContext: QueryConte
Name.PROJECT -> emitProjectStage(stage)
else -> {}
}
emitObjectValueEnd()
}
emitArrayEnd(long = true)
return this
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
package com.mongodb.jbplugin.dialects.mongosh.aggr

import com.mongodb.jbplugin.dialects.mongosh.assertGeneratedQuery
import com.mongodb.jbplugin.mql.BsonInt32
import com.mongodb.jbplugin.mql.BsonString
import com.mongodb.jbplugin.mql.Node
import com.mongodb.jbplugin.mql.QueryContext
import com.mongodb.jbplugin.mql.components.HasAggregation
import com.mongodb.jbplugin.mql.components.HasFieldReference
import com.mongodb.jbplugin.mql.components.HasFilter
import com.mongodb.jbplugin.mql.components.HasProjections
import com.mongodb.jbplugin.mql.components.HasValueReference
import com.mongodb.jbplugin.mql.components.IsCommand
import com.mongodb.jbplugin.mql.components.Name
Expand All @@ -15,14 +17,15 @@ import org.junit.jupiter.api.Test

class AggregateTest {
@Test
fun `can format an aggregate query with a match expression at the beginning`() {
fun `can format a safe explain command for a valid aggregate query`() {
assertGeneratedQuery(
"""
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}match": {"myField": "myVal"}}])
""".trimIndent()
db.getSiblingDB(database).getCollection(collection).explain("queryPlanner").aggregate([{"${'$'}match": {"myField": "myVal"}}, ])
""".trimIndent(),
explain = QueryContext.ExplainPlanType.SAFE
) {
Node(
Unit,
Expand Down Expand Up @@ -66,15 +69,15 @@ class AggregateTest {
}

@Test
fun `can format a safe explain command for a valid aggregate query`() {
fun `can format a full explain command for a valid aggregate query`() {
assertGeneratedQuery(
"""
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).explain("queryPlanner").aggregate([{"${'$'}match": {"myField": "myVal"}}])
db.getSiblingDB(database).getCollection(collection).explain("executionStats").aggregate([{"${'$'}match": {"myField": "myVal"}}, ])
""".trimIndent(),
explain = QueryContext.ExplainPlanType.SAFE
explain = QueryContext.ExplainPlanType.FULL
) {
Node(
Unit,
Expand Down Expand Up @@ -118,13 +121,93 @@ class AggregateTest {
}

@Test
fun `can format a full explain command for a valid aggregate query`() {
fun `can format an aggregate with multiple stages`() {
assertGeneratedQuery(
"""
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).aggregate([{"${'$'}match": {"myField": "myVal"}}, {"${'$'}project": {"myField": 1, }}, ])
""".trimIndent(),
explain = QueryContext.ExplainPlanType.NONE
) {
Node(
Unit,
listOf(
IsCommand(IsCommand.CommandType.AGGREGATE),
HasAggregation(
listOf(
Node(
Unit,
listOf(
Named(Name.MATCH),
HasFilter(
listOf(
Node(
Unit,
listOf(
HasFieldReference(
HasFieldReference.FromSchema(
Unit,
"myField"
)
),
HasValueReference(
HasValueReference.Constant(
Unit,
"myVal",
BsonString
)
)
)
)
)
),
)
),
Node(
Unit,
listOf(
Named(Name.PROJECT),
HasProjections(
listOf(
Node(
Unit,
listOf(
HasFieldReference(
HasFieldReference.FromSchema(
Unit,
"myField"
)
),
HasValueReference(
HasValueReference.Constant(
Unit,
1,
BsonInt32
)
)
)
)
)
),
)
)
)
)
)
)
}
}

@Test
fun `can format a full explain command for a valid aggregate query removing destructive operations`() {
assertGeneratedQuery(
"""
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).explain("executionStats").aggregate([{"${'$'}match": {"myField": "myVal"}}])
db.getSiblingDB(database).getCollection(collection).explain("executionStats").aggregate([{"${'$'}match": {"myField": "myVal"}}, ])
""".trimIndent(),
explain = QueryContext.ExplainPlanType.FULL
) {
Expand Down Expand Up @@ -159,6 +242,29 @@ class AggregateTest {
)
)
)
),
Named(Name.PROJECT),
HasProjections(
listOf(
Node(
Unit,
listOf(
HasFieldReference(
HasFieldReference.FromSchema(
Unit,
"myField"
)
),
HasValueReference(
HasValueReference.Constant(
Unit,
"myVal",
BsonString
)
)
)
)
)
)
)
)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
package com.mongodb.jbplugin.dialects.mongosh.aggr

import com.mongodb.jbplugin.dialects.mongosh.assertGeneratedQuery
import com.mongodb.jbplugin.mql.BsonString
import com.mongodb.jbplugin.mql.Node
import com.mongodb.jbplugin.mql.components.HasAggregation
import com.mongodb.jbplugin.mql.components.HasFieldReference
import com.mongodb.jbplugin.mql.components.HasFilter
import com.mongodb.jbplugin.mql.components.HasValueReference
import com.mongodb.jbplugin.mql.components.IsCommand
import com.mongodb.jbplugin.mql.components.Name
import com.mongodb.jbplugin.mql.components.Named
import org.junit.jupiter.api.Test

class MatchTest {
@Test
fun `can format an aggregate query with a match expression at the beginning`() {
assertGeneratedQuery(
"""
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}match": {"myField": "myVal"}}, ])
""".trimIndent()
) {
Node(
Unit,
listOf(
IsCommand(IsCommand.CommandType.AGGREGATE),
HasAggregation(
listOf(
Node(
Unit,
listOf(
Named(Name.MATCH),
HasFilter(
listOf(
Node(
Unit,
listOf(
HasFieldReference(
HasFieldReference.FromSchema(
Unit,
"myField"
)
),
HasValueReference(
HasValueReference.Constant(
Unit,
"myVal",
BsonString
)
)
)
)
)
)
)
)
)
)
)
)
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ class ProjectTest {
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}project": {"myField": 1, }}])
db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}project": {"myField": 1, }}, ])
""".trimIndent()
) {
Node(
Expand Down Expand Up @@ -71,7 +71,7 @@ class ProjectTest {
var collection = ""
var database = ""
db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}project": {"myField": -1, }}])
db.getSiblingDB(database).getCollection(collection).aggregate([{"${"$"}project": {"myField": -1, }}, ])
""".trimIndent()
) {
Node(
Expand Down

0 comments on commit c40b321

Please sign in to comment.