feat(query-generation): modal for specifying default values INTELLIJ-198

CHANGELOG.md

@@ -5,6 +5,7 @@ MongoDB plugin for IntelliJ IDEA.
 ## [Unreleased]
 
 ### Added
+* [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.
 * [INTELLIJ-195](https://jira.mongodb.org/browse/INTELLIJ-195) Add support for $unwind when generating the query into DataGrip.

buildSrc/src/main/kotlin/com.mongodb.intellij.plugin-component.gradle.kts

@@ -64,6 +64,7 @@ dependencies {
     testImplementation(libs.testing.spring.mongodb)
     testImplementation(libs.testing.jsoup)
     testImplementation(libs.testing.video.recorder)
+    testImplementation(libs.testing.assertj.swing)
     testImplementation(libs.testing.remoteRobot)
     testImplementation(libs.testing.remoteRobotDeps.remoteFixtures)
     testImplementation(libs.testing.remoteRobotDeps.ideLauncher)

gradle/libs.versions.toml

@@ -34,6 +34,7 @@ testContainers = "1.19.8"
 spring-mongodb="4.3.2"
 semver-parser="2.0.0"
 snakeyaml="2.3"
+assertj-swing="3.17.1"
 
 [libraries]
 ## Kotlin compileOnly libraries. They must not be bundled because they are already part of the
@@ -75,7 +76,7 @@ testing-testContainers-core = { group = "org.testcontainers", name = "testcontai
 testing-testContainers-mongodb = { group = "org.testcontainers", name = "mongodb", version.ref = "testContainers" }
 testing-testContainers-jupiter = { group = "org.testcontainers", name = "junit-jupiter", version.ref = "testContainers" }
 testing-spring-mongodb = { group = "org.springframework.data", name="spring-data-mongodb", version.ref="spring-mongodb" }
-
+testing-assertj-swing = { group = "org.assertj", name="assertj-swing", version.ref="assertj-swing" }
 ######################################################
 ## Libraries and plugins only used for the buildScript.
 buildScript-plugin-kotlin = { group = "org.jetbrains.kotlin", name = "kotlin-gradle-plugin", version.ref = "kotlin-stdlib" }

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/codeActions/impl/RunQueryCodeActionBridge.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.PsiLiteralExpression
 import com.mongodb.jbplugin.accessadapter.datagrip.adapter.isConnected
 import com.mongodb.jbplugin.codeActions.AbstractMongoDbCodeActionBridge
 import com.mongodb.jbplugin.codeActions.MongoDbCodeAction
+import com.mongodb.jbplugin.codeActions.impl.runQuery.RunQueryModal
 import com.mongodb.jbplugin.codeActions.sourceForMarker
 import com.mongodb.jbplugin.dialects.DialectFormatter
 import com.mongodb.jbplugin.dialects.OutputQuery
@@ -35,7 +36,6 @@ import com.mongodb.jbplugin.i18n.CodeActionsMessages
 import com.mongodb.jbplugin.i18n.Icons
 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.observability.TelemetryEvent
 import com.mongodb.jbplugin.observability.TelemetryEvent.QueryRunEvent.Console
@@ -74,23 +74,34 @@ internal object RunQueryCodeAction : MongoDbCodeAction {
                 } else {
                     emitRunQueryEvent(query, dataSource)
 
-                    coroutineScope.launchChildBackground {
-                        val outputQuery = MongoshDialect.formatter.formatQuery(
-                            query,
-                            QueryContext.empty(prettyPrint = true)
-                        )
+                    if (dataSource == null || dataSource.isConnected() == false) {
+                        return@LineMarkerInfo
+                    }
 
-                        if (dataSource?.isConnected() == true) {
-                            coroutineScope.launchChildOnUi {
-                                openDataGripConsole(query, dataSource, outputQuery.query)
-                            }
-                        } else {
-                            openConsoleAfterSelection(
+                    val queryContext = RunQueryModal(
+                        query,
+                        dataSource,
+                        coroutineScope
+                    ).askForQueryContext()
+                    if (queryContext != null) {
+                        coroutineScope.launchChildBackground {
+                            val outputQuery = MongoshDialect.formatter.formatQuery(
                                 query,
-                                outputQuery,
-                                query.source.project,
-                                coroutineScope
+                                queryContext
                             )
+
+                            if (dataSource.isConnected() == true) {
+                                coroutineScope.launchChildOnUi {
+                                    openDataGripConsole(query, dataSource, outputQuery.query)
+                                }
+                            } else {
+                                openConsoleAfterSelection(
+                                    query,
+                                    outputQuery,
+                                    query.source.project,
+                                    coroutineScope
+                                )
+                            }
                         }
                     }
                 }

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/codeActions/impl/runQuery/NamespaceSelector.kt

@@ -0,0 +1,179 @@
+package com.mongodb.jbplugin.codeActions.impl.runQuery
+
+import com.intellij.collaboration.ui.selectFirst
+import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.rd.util.launchChildBackground
+import com.intellij.openapi.ui.ComboBox
+import com.intellij.ui.AnimatedIcon.ANIMATION_IN_RENDERER_ALLOWED
+import com.intellij.ui.components.JBLabel
+import com.mongodb.jbplugin.accessadapter.datagrip.DataGripBasedReadModelProvider
+import com.mongodb.jbplugin.accessadapter.slice.ListCollections
+import com.mongodb.jbplugin.accessadapter.slice.ListDatabases
+import com.mongodb.jbplugin.i18n.Icons
+import com.mongodb.jbplugin.i18n.Icons.scaledToText
+import com.mongodb.jbplugin.meta.latest
+import com.mongodb.jbplugin.meta.service
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.MutableSharedFlow
+import java.awt.Component
+import javax.swing.DefaultComboBoxModel
+import javax.swing.SwingConstants
+
+class NamespaceSelector(
+    private val project: Project,
+    private val dataSource: LocalDataSource,
+    private val coroutineScope: CoroutineScope,
+) {
+    sealed interface Event {
+        data object DatabasesLoading : Event
+        data class DatabasesLoaded(val databases: List<String>) : Event
+        data class DatabaseSelected(val database: String) : Event
+        data object CollectionsLoading : Event
+        data class CollectionsLoaded(val collections: List<String>) : Event
+        data class CollectionSelected(val collection: String) : Event
+    }
+
+    private val events: MutableSharedFlow<Event> = MutableSharedFlow(extraBufferCapacity = 1)
+
+    private val databaseModel = DefaultComboBoxModel<String>(emptyArray())
+    val databaseComboBox = ComboBox<String?>(databaseModel)
+    private val collectionModel = DefaultComboBoxModel<String>(emptyArray())
+    val collectionComboBox = ComboBox<String?>(collectionModel)
+
+    val selectedDatabase: String?
+        get() = databaseModel.selectedItem?.toString()
+
+    val selectedCollection: String?
+        get() = collectionModel.selectedItem?.toString()
+
+    private val loadingDatabases: Boolean by events.latest(
+        onNewEvent = { event, state ->
+            when (event) {
+                is Event.DatabasesLoading -> true
+                is Event.DatabasesLoaded -> false
+                else -> state
+            }
+        },
+        onChange = {},
+        defaultValue = true
+    )
+
+    private val loadingCollections: Boolean by events.latest(
+        onNewEvent = { event, state ->
+            when (event) {
+                is Event.CollectionsLoading -> true
+                is Event.CollectionsLoaded -> false
+                else -> state
+            }
+        },
+        onChange = {},
+        defaultValue = false
+    )
+
+    init {
+        databaseComboBox.isEnabled = false
+        collectionComboBox.isEnabled = false
+        databaseComboBox.name = "DatabaseComboBox"
+        collectionComboBox.name = "CollectionComboBox"
+
+        databaseComboBox.prototypeDisplayValue = "XXXXXXXXXXXXXXXXXXXXX"
+        databaseComboBox.putClientProperty(ANIMATION_IN_RENDERER_ALLOWED, true)
+
+        collectionComboBox.prototypeDisplayValue = "XXXXXXXXXXXXXXXXXXXXX"
+        collectionComboBox.putClientProperty(ANIMATION_IN_RENDERER_ALLOWED, true)
+
+        databaseComboBox.setRenderer { _, value, index, _, _ -> renderDatabaseItem(value, index) }
+        collectionComboBox.setRenderer { _, value, index, _, _ ->
+            renderCollectionItem(value, index)
+        }
+
+        coroutineScope.launchChildBackground {
+            events.collect(::handleEvent)
+        }
+
+        databaseComboBox.addItemListener {
+            coroutineScope.launchChildBackground {
+                events.emit(Event.DatabaseSelected(it.item.toString()))
+            }
+        }
+
+        collectionComboBox.addItemListener {
+            coroutineScope.launchChildBackground {
+                events.emit(Event.CollectionSelected(it.item.toString()))
+            }
+        }
+
+        loadDatabases()
+    }
+
+    private fun loadDatabases() {
+        coroutineScope.launchChildBackground {
+            events.emit(Event.DatabasesLoading)
+            val readModel by project.service<DataGripBasedReadModelProvider>()
+            val result = readModel.slice(
+                dataSource,
+                ListDatabases.Slice
+            )
+
+            events.emit(Event.DatabasesLoaded(result.databases.map { it.name }))
+        }
+    }
+
+    private fun handleEvent(event: Event) {
+        when (event) {
+            is Event.DatabasesLoading -> {
+                databaseComboBox.isEnabled = false
+            }
+
+            is Event.DatabasesLoaded -> {
+                databaseModel.removeAllElements()
+                collectionModel.removeAllElements()
+                databaseModel.addAll(event.databases)
+                databaseModel.selectFirst()
+                databaseComboBox.isEnabled = true
+            }
+            is Event.DatabaseSelected -> {
+                coroutineScope.launchChildBackground {
+                    events.emit(Event.CollectionsLoading)
+                    val readModel by project.service<DataGripBasedReadModelProvider>()
+                    val result = readModel.slice(
+                        dataSource,
+                        ListCollections.Slice(event.database)
+                    )
+
+                    events.emit(Event.CollectionsLoaded(result.collections.map { it.name }))
+                }
+            }
+            is Event.CollectionsLoading -> {
+                collectionModel.removeAllElements()
+                collectionComboBox.isEnabled = false
+            }
+
+            is Event.CollectionsLoaded -> {
+                collectionModel.addAll(event.collections)
+                collectionModel.selectFirst()
+                collectionComboBox.isEnabled = true
+            }
+            else -> {}
+        }
+    }
+
+    private fun renderDatabaseItem(item: String?, index: Int): Component = if (item == null &&
+        index == -1 &&
+        loadingDatabases
+    ) {
+        JBLabel("Loading databases...", Icons.loading.scaledToText(), SwingConstants.LEFT)
+    } else {
+        JBLabel(item ?: "", Icons.databaseAutocompleteEntry, SwingConstants.LEFT)
+    }
+
+    private fun renderCollectionItem(item: String?, index: Int): Component = if (item == null &&
+        index == -1 &&
+        loadingCollections
+    ) {
+        JBLabel("Loading collections...", Icons.loading.scaledToText(), SwingConstants.LEFT)
+    } else {
+        JBLabel(item ?: "", Icons.collectionAutocompleteEntry, SwingConstants.LEFT)
+    }
+}