Skip to content

Commit

Permalink
chore: now missing tests and small fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
kmruiz committed Nov 22, 2024
1 parent 8f1391a commit 6c65c9f
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -54,10 +54,13 @@ public List<Document> queryMoviesByYear(String year) {
.aggregate(
List.of(
Aggregates.match(
Filters.eq("_id", year)
Filters.eq("year", year)
),
Aggregates.group(
"newField", Accumulators.avg("test", "$year")
"newField",
Accumulators.avg("test", "$xxx"),
Accumulators.sum("test2", "$year"),
Accumulators.bottom("field", Sorts.ascending("year"), "$year")
),
Aggregates.project(
Projections.fields(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ private const val ACCUMULATORS_FQN = "com.mongodb.client.model.Accumulators"
private const val SORTS_FQN = "com.mongodb.client.model.Sorts"
private const val JAVA_LIST_FQN = "java.util.List"
private const val JAVA_ARRAYS_FQN = "java.util.Arrays"

private val PARSEABLE_AGGREGATION_STAGE_METHODS = listOf(
"match",
"project",
Expand Down Expand Up @@ -228,7 +229,8 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {

return containingClass.qualifiedName == AGGREGATES_FQN ||
containingClass.qualifiedName == PROJECTIONS_FQN ||
containingClass.qualifiedName == SORTS_FQN
containingClass.qualifiedName == SORTS_FQN ||
containingClass.qualifiedName == ACCUMULATORS_FQN
}

private fun resolveBsonBuilderCall(
Expand Down Expand Up @@ -297,7 +299,7 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
}
val valueExpression = filter.argumentList.expressions[0]
val valueReference = resolveValueFromExpression(valueExpression)
val fieldReference = HasFieldReference.FromSchema(valueExpression, "_id")
val fieldReference = FromSchema(valueExpression, "_id")

return Node(
filter,
Expand Down Expand Up @@ -505,7 +507,7 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
}

val accumulators = stageCall.getVarArgsOrIterableArgs().drop(1)
.mapNotNull { resolveToAccumulatorCall(it) }
.mapNotNull { resolveBsonBuilderCall(it, ACCUMULATORS_FQN) }

val parsedAccumulators = accumulators.mapNotNull { parseAccumulatorExpression(it) }
return nodeWithAccumulators(parsedAccumulators)
Expand All @@ -515,13 +517,10 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
}
}

private fun resolveToProjectionCall(element: PsiElement): PsiMethodCallExpression? {
return element.resolveToMethodCallExpression { _, methodCall ->
methodCall.containingClass?.qualifiedName == PROJECTIONS_FQN
}
}

private fun parseProjectionExpression(expression: PsiMethodCallExpression): List<Node<PsiElement>> {
private fun parseBsonBuilderCallsSimilarToProjections(
expression: PsiMethodCallExpression,
classQualifiedName: String
): List<Node<PsiElement>> {
val methodCall = expression.resolveMethod() ?: return emptyList()
return when (methodCall.name) {
"fields",
Expand Down Expand Up @@ -556,7 +555,8 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
)
)
)
) else -> null
)
else -> null
}
}

Expand All @@ -572,8 +572,8 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
"avg" -> parseKeyValAccumulator(expression, Name.AVG)
"first" -> parseKeyValAccumulator(expression, Name.FIRST)
"last" -> parseKeyValAccumulator(expression, Name.LAST)
"top" -> null
"bottom" -> null
"top" -> parseLeadingAccumulatorExpression(expression, Name.TOP)
"bottom" -> parseLeadingAccumulatorExpression(expression, Name.BOTTOM)
"max" -> parseKeyValAccumulator(expression, Name.MAX)
"min" -> parseKeyValAccumulator(expression, Name.MIN)
"push" -> parseKeyValAccumulator(expression, Name.PUSH)
Expand Down Expand Up @@ -605,6 +605,34 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
)
}

private fun parseLeadingAccumulatorExpression(expression: PsiMethodCallExpression, name: Name): Node<PsiElement>? {
val keyExpr = expression.argumentList.expressions.getOrNull(0) ?: return null
val sortExprArgument = expression.argumentList.expressions.getOrNull(1) ?: return null
val valueExpr = expression.argumentList.expressions.getOrNull(2) ?: return null

val sortExpr = resolveBsonBuilderCall(sortExprArgument, SORTS_FQN) ?: return null

val fieldName = keyExpr.tryToResolveAsConstantString()
val sort = parseBsonBuilderCallsSimilarToProjections(sortExpr, SORTS_FQN)
val accumulatorExpr = parseComputedExpression(valueExpr, createsNewField = false)

return Node(
expression,
listOf(
Named(name),
HasFieldReference(
if (fieldName != null) {
Computed(keyExpr, fieldName, fieldName)
} else {
HasFieldReference.Unknown as HasFieldReference.FieldReference<PsiElement>
}
),
HasSorts(sort),
accumulatorExpr
)
)
}

private fun hasMongoDbSessionReference(methodCall: PsiMethodCallExpression): Boolean {
val hasEnoughArgs = methodCall.argumentList.expressionCount > 0
if (!hasEnoughArgs) {
Expand All @@ -615,6 +643,31 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
return typeOfFirstArg != null && typeOfFirstArg.equalsToText(SESSION_FQN)
}

private fun parseComputedExpression(element: PsiElement, createsNewField: Boolean = true): HasValueReference<PsiElement> {
return HasValueReference(
when (val expression = element.tryToResolveAsConstantString()) {
null -> HasValueReference.Unknown as HasValueReference.ValueReference<PsiElement>
else -> HasValueReference.Computed(
element,
Node(
element,
listOf(
if (createsNewField) {
HasFieldReference(
Computed(element, expression.trim('$'), expression)
)
} else {
HasFieldReference(
FromSchema(element, expression.trim('$'), expression)
)
}
)
)
)
}
)
}

private fun isAggregationStageMethodCall(callMethod: PsiMethod?): Boolean {
return PARSEABLE_AGGREGATION_STAGE_METHODS.contains(callMethod?.name) &&
callMethod?.containingClass?.qualifiedName == AGGREGATES_FQN
Expand All @@ -624,7 +677,7 @@ object JavaDriverDialectParser : DialectParser<PsiElement> {
val fieldNameAsString = expression.tryToResolveAsConstantString()
val fieldReference =
fieldNameAsString?.let {
HasFieldReference.FromSchema(expression, it)
FromSchema(expression, it)
} ?: HasFieldReference.Unknown

return fieldReference
Expand Down

0 comments on commit 6c65c9f

Please sign in to comment.