chore: draft toolbar

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/editor/EditorToolbarDecorator.kt

@@ -1,24 +1,132 @@
 package com.mongodb.jbplugin.editor
 
+import com.intellij.database.dataSource.DatabaseConnection
+import com.intellij.database.dataSource.DatabaseConnectionManager
+import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.database.dataSource.LocalDataSourceManager
+import com.intellij.database.dataSource.connection.ConnectionRequestor
+import com.intellij.database.model.RawDataSource
+import com.intellij.database.psi.DataSourceManager
+import com.intellij.database.run.ConsoleRunConfiguration
 import com.intellij.lang.java.JavaLanguage
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.event.EditorFactoryEvent
 import com.intellij.openapi.editor.event.EditorFactoryListener
 import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.openapi.rd.util.launchChildOnUi
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.removeUserData
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiManager
+import com.intellij.util.messages.MessageBusConnection
+import com.mongodb.jbplugin.accessadapter.datagrip.adapter.isConnected
+import com.mongodb.jbplugin.accessadapter.datagrip.adapter.isMongoDbDataSource
+import kotlinx.coroutines.*
 
 /**
  * This decorator listens to an IntelliJ Editor lifecycle
  * and attaches our toolbar if necessary.
  */
-class EditorToolbarDecorator : EditorFactoryListener {
-    private val toolbar = MdbJavaEditorToolbar()
+class EditorToolbarDecorator(
+    private val coroutineScope: CoroutineScope,
+) : EditorFactoryListener,
+    DataSourceManager.Listener,
+    DatabaseConnectionManager.Listener {
+    /**
+     * This needs to be synchronised with the MongoDbVirtualFileDataSourceProvider field with the same name.
+     *
+     * @see MongoDbVirtualFileDataSourceProvider
+     */
+    private val attachedDataSource: Key<LocalDataSource> = Key.create("com.mongodb.jbplugin.AttachedDataSource")
+    private val toolbar =
+        MdbJavaEditorToolbar(
+            onDataSourceSelected = this::onDataSourceSelected,
+            onDataSourceUnselected = this::onDataSourceUnselected,
+        )
+
     private lateinit var editor: Editor
+    private lateinit var connection: MessageBusConnection
+
+    private fun onDataSourceSelected(dataSource: LocalDataSource) {
+        editor.putUserData(attachedDataSource, dataSource)
+        val project = editor.project ?: return
+
+        if (!dataSource.isConnected()) {
+            coroutineScope.launch {
+                /**
+                 * We don't track connection failures here, as we really don't care. If we have a connection, attach
+                 * it, if we don't, reset the form. We already have ConnectionFailureProbe for error tracking.
+                 *
+                 * @see com.mongodb.jbplugin.observability.probe.ConnectionFailureProbe
+                 */
+                toolbar.connecting = true
+                val connectionManager = DatabaseConnectionManager.getInstance()
+                val connectionHandler =
+                    connectionManager
+                        .build(project, dataSource)
+                        .setRequestor(ConnectionRequestor.Anonymous())
+                        .setAskPassword(true)
+                        .setRunConfiguration(
+                            ConsoleRunConfiguration.newConfiguration(project).apply {
+                                setOptionsFromDataSource(dataSource)
+                            },
+                        )
+
+                val connection =
+                    async {
+                        val loadingAnimation =
+                            launchChildOnUi {
+                                while (true) {
+                                    delay(50)
+                                    toolbar.updateUI()
+                                }
+                            }
+                        val connectionJob = runCatching { connectionHandler.create()?.get() }
+                        connectionJob.onFailure {
+                            toolbar.connecting = false
+                            toolbar.failedConnection = dataSource
+                        }
+
+                        loadingAnimation.cancelAndJoin()
+                        connectionJob.getOrNull()
+                    }.await()
+
+                if (toolbar.failedConnection != null) {
+                    return@launch
+                }
+
+                toolbar.connecting = false
+                toolbar.updateUI()
+
+                if (connection == null || !dataSource.isConnected()) { // could not connect, do nothing
+                    toolbar.selectedDataSource = null // remove data source because we didn't connect
+                    return@launch
+                }
+
+                editor.virtualFile?.putUserData(attachedDataSource, dataSource)
+            }
+        } else {
+            editor.virtualFile?.putUserData(attachedDataSource, dataSource)
+        }
+    }
+
+    private fun onDataSourceUnselected() {
+        editor.virtualFile?.removeUserData(attachedDataSource)
+    }
 
     override fun editorCreated(event: EditorFactoryEvent) {
         editor = event.editor
 
+        if (editor.project != null) {
+            val project = editor.project!!
+            connection = project.messageBus.connect()
+            connection.subscribe(DataSourceManager.TOPIC, this)
+            connection.subscribe(DatabaseConnectionManager.TOPIC, this)
+
+            val localDataSourceManager = DataSourceManager.byDataSource(project, LocalDataSource::class.java) ?: return
+            toolbar.dataSources = localDataSourceManager.dataSources.filter { it.isMongoDbDataSource() }
+        }
+
         ensureToolbarIsVisibleIfNecessary()
     }
 
@@ -58,4 +166,50 @@ class EditorToolbarDecorator : EditorFactoryListener {
             return@any it.importReference?.canonicalText?.startsWith("com.mongodb") == true
         }
     }
+
+    override fun <T : RawDataSource?> dataSourceAdded(
+        manager: DataSourceManager<T>,
+        dataSource: T & Any,
+    ) {
+        val localDataSourceManager = manager as? LocalDataSourceManager ?: return
+        toolbar.dataSources = localDataSourceManager.dataSources.filter { it.isMongoDbDataSource() }
+    }
+
+    override fun <T : RawDataSource?> dataSourceRemoved(
+        manager: DataSourceManager<T>,
+        dataSource: T & Any,
+    ) {
+        val localDataSourceManager = manager as? LocalDataSourceManager ?: return
+
+        if (toolbar.selectedDataSource?.uniqueId == dataSource.uniqueId) {
+            toolbar.selectedDataSource = null
+        }
+
+        toolbar.dataSources = localDataSourceManager.dataSources.filter { it.isMongoDbDataSource() }
+    }
+
+    override fun <T : RawDataSource?> dataSourceChanged(
+        manager: DataSourceManager<T>?,
+        dataSource: T?,
+    ) {
+        val localDataSourceManager = manager as? LocalDataSourceManager ?: return
+
+        if (toolbar.selectedDataSource?.uniqueId == dataSource?.uniqueId) {
+            toolbar.selectedDataSource = null
+        }
+
+        toolbar.dataSources = localDataSourceManager.dataSources.filter { it.isMongoDbDataSource() }
+    }
+
+    override fun connectionChanged(
+        connection: DatabaseConnection,
+        added: Boolean,
+    ) {
+        val dataSource = connection.connectionPoint.dataSource
+        val selectedDataSource = toolbar.selectedDataSource
+
+        if (dataSource.isMongoDbDataSource() && !dataSource.isConnected() && selectedDataSource?.uniqueId == dataSource.uniqueId) {
+            toolbar.selectedDataSource = null
+        }
+    }
 }

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/editor/MdbJavaEditorToolbar.kt

@@ -1,14 +1,116 @@
 package com.mongodb.jbplugin.editor
 
+import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.openapi.ui.ComboBox
+import com.intellij.openapi.ui.asSequence
+import com.intellij.sql.indexOf
+import com.intellij.ui.AnimatedIcon.ANIMATION_IN_RENDERER_ALLOWED
 import com.intellij.ui.components.JBLabel
 import com.intellij.ui.components.JBPanel
+import com.mongodb.jbplugin.accessadapter.datagrip.adapter.isConnected
+import com.mongodb.jbplugin.i18n.Icons
+import com.mongodb.jbplugin.i18n.Icons.scaledToText
+import com.mongodb.jbplugin.i18n.MdbToolbarMessages
 import java.awt.BorderLayout
+import java.awt.event.ItemEvent.DESELECTED
+import javax.swing.DefaultComboBoxModel
+import javax.swing.SwingConstants
 
 /**
  * Represents the toolbar that will be inserted into an active Java editor.
  */
-class MdbJavaEditorToolbar : JBPanel<MdbJavaEditorToolbar>(BorderLayout()) {
+class MdbJavaEditorToolbar(
+    private val onDataSourceSelected: DataSourceSelectedListener,
+    private val onDataSourceUnselected: DataSourceUnselectedListener,
+) : JBPanel<MdbJavaEditorToolbar>(BorderLayout()) {
+    var connecting: Boolean = false
+    var failedConnection: LocalDataSource? = null
+
+    private val connectionComboBox =
+        ComboBox(
+            DefaultComboBoxModel(
+                emptyArray<LocalDataSource>(),
+            ),
+        )
+
     init {
-        add(JBLabel("Placeholder"))
+        add(connectionComboBox, BorderLayout.EAST)
+        connectionComboBox.addItemListener {
+            failedConnection = null
+            connecting = false
+
+            if (it.stateChange == DESELECTED) {
+                onDataSourceUnselected()
+            } else {
+                onDataSourceSelected(it.item as LocalDataSource)
+            }
+        }
+        connectionComboBox.putClientProperty(ANIMATION_IN_RENDERER_ALLOWED, true)
+        connectionComboBox.setRenderer { _, value, index, _, _ ->
+            if (value == null && index == -1) {
+                JBLabel(MdbToolbarMessages.message("attach.datasource.to.editor"), Icons.Logo.scaledToText(), SwingConstants.LEFT)
+            } else if (value == null) {
+                JBLabel(MdbToolbarMessages.message("detach.datasource.from.editor"), Icons.Remove.scaledToText(), SwingConstants.LEFT)
+            } else {
+                val icon =
+                    if (value.isConnected()) {
+                        Icons.LogoConnected.scaledToText()
+                    } else if (connecting) {
+                        Icons.LoadingIcon.scaledToText()
+                    } else if (failedConnection?.uniqueId == value.uniqueId) {
+                        Icons.ConnectionFailed.scaledToText()
+                    } else {
+                        Icons.Logo.scaledToText()
+                    }
+                JBLabel(value.name, icon, SwingConstants.LEFT)
+            }
+        }
+    }
+
+    var dataSources: List<LocalDataSource>
+        set(value) {
+            failedConnection = null
+            connecting = false
+
+            val selectedItem = this.selectedDataSource?.uniqueId
+            val model = connectionComboBox.model as DefaultComboBoxModel<LocalDataSource>
+            model.removeAllElements()
+            model.addElement(null)
+            model.addAll(value)
+            selectDataSourceWithId(selectedItem)
+        }
+        get() = (connectionComboBox.model as DefaultComboBoxModel<LocalDataSource>).asSequence().toList().filterNotNull()
+
+    var selectedDataSource: LocalDataSource?
+        set(value) {
+            failedConnection = null
+            connecting = false
+
+            if (value == null) {
+                connectionComboBox.selectedItem = null
+            } else {
+                connectionComboBox.selectedItem = value
+            }
+        }
+        get() = connectionComboBox.selectedItem as? LocalDataSource
+
+    private fun selectDataSourceWithId(id: String?) {
+        if (id == null) {
+            connectionComboBox.selectedItem = null
+        } else {
+            val dataSourceIndex =
+                (connectionComboBox.model as DefaultComboBoxModel<LocalDataSource?>)
+                    .asSequence()
+                    .toList()
+                    .indexOf { it?.uniqueId == id }
+            if (dataSourceIndex == -1) {
+                connectionComboBox.selectedItem = null
+            } else {
+                connectionComboBox.selectedIndex = dataSourceIndex
+            }
+        }
     }
 }
+
+typealias DataSourceSelectedListener = (LocalDataSource) -> Unit
+typealias DataSourceUnselectedListener = () -> Unit

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/editor/MongoDbVirtualFileDataSourceProvider.kt

@@ -0,0 +1,28 @@
+package com.mongodb.jbplugin.editor
+
+import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.database.psi.DbDataSource
+import com.intellij.database.psi.DbPsiFacade
+import com.intellij.database.util.VirtualFileDataSourceProvider
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.vfs.VirtualFile
+
+class MongoDbVirtualFileDataSourceProvider : VirtualFileDataSourceProvider() {
+    /**
+     * This needs to be synchronised with the EditorToolbarDecorator field with the same name.
+     *
+     * @see EditorToolbarDecorator
+     */
+    private val attachedDataSource: Key<LocalDataSource> = Key.create("com.mongodb.jbplugin.AttachedDataSource")
+
+    override fun getDataSource(
+        project: Project,
+        file: VirtualFile,
+    ): DbDataSource? {
+        val facade = DbPsiFacade.getInstance(project)
+        val attachedDataSource = file.getUserData(attachedDataSource) ?: return null
+
+        return facade.findDataSource(attachedDataSource.uniqueId)
+    }
+}

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/i18n/Icons.kt

@@ -0,0 +1,32 @@
+package com.mongodb.jbplugin.i18n
+
+import com.intellij.icons.AllIcons
+import com.intellij.ide.ui.NotRoamableUiSettings
+import com.intellij.openapi.util.IconLoader
+import com.intellij.ui.AnimatedIcon
+import com.intellij.ui.LayeredIcon
+import com.intellij.util.IconUtil
+import java.awt.*
+import javax.swing.Icon
+import javax.swing.SwingConstants
+
+object Icons {
+    val LoadingIcon = AnimatedIcon.Default()
+    val GreenCircle = IconLoader.getIcon("/icons/GreenCircle.svg", javaClass)
+
+    val Logo = AllIcons.Providers.MongoDB
+    val LogoConnected =
+        LayeredIcon.layeredIcon(arrayOf(Logo, GreenCircle)).apply {
+            val scaledGreenCircle = IconUtil.resizeSquared(GreenCircle, 6)
+            setIcon(scaledGreenCircle, 1, SwingConstants.SOUTH_EAST)
+        }
+
+    val ConnectionFailed = AllIcons.General.Error
+    val Remove = AllIcons.Diff.Remove
+
+    fun Icon.scaledToText(parentComponent: Component? = null): Icon {
+        val settingsManager: NotRoamableUiSettings = NotRoamableUiSettings.getInstance()
+        val settings = settingsManager.state
+        return IconUtil.scaleByFont(this, parentComponent, settings.fontSize)
+    }
+}

packages/jetbrains-plugin/src/main/kotlin/com/mongodb/jbplugin/i18n/MdbToolbarMessages.kt

@@ -0,0 +1,19 @@
+package com.mongodb.jbplugin.i18n
+
+import com.intellij.DynamicBundle
+import org.jetbrains.annotations.Nls
+import org.jetbrains.annotations.NonNls
+import org.jetbrains.annotations.PropertyKey
+
+object MdbToolbarMessages {
+    @NonNls
+    private const val BUNDLE = "messages.MdbToolbarBundle"
+    private val instance = DynamicBundle(MdbToolbarMessages::class.java, BUNDLE)
+
+    fun message(
+        key:
+            @PropertyKey(resourceBundle = BUNDLE)
+            String,
+        vararg params: Any,
+    ): @Nls String = instance.getMessage(key, *params)
+}

packages/jetbrains-plugin/src/main/resources/messages/MdbToolbarBundle.properties

@@ -0,0 +1,2 @@
+attach.datasource.to.editor=Attach MongoDB data source
+detach.datasource.from.editor=Detach data source