Remove usages of deprecated API in DatabaseAuthProvider

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/datagrip/auth/SecretsManagerAuth.kt

@@ -7,12 +7,10 @@ import com.fasterxml.jackson.databind.DeserializationFeature
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.fasterxml.jackson.module.kotlin.readValue
 import com.intellij.credentialStore.Credentials
-import com.intellij.database.access.DatabaseCredentials
 import com.intellij.database.dataSource.DatabaseAuthProvider.AuthWidget
 import com.intellij.database.dataSource.DatabaseConnectionInterceptor.ProtoConnection
 import com.intellij.database.dataSource.DatabaseCredentialsAuthProvider
 import com.intellij.database.dataSource.LocalDataSource
-import com.intellij.openapi.project.Project
 import kotlinx.coroutines.future.future
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient
 import software.aws.toolkits.core.ConnectionSettings
@@ -45,8 +43,7 @@ class SecretsManagerAuth : DatabaseAuthProviderCompatabilityAdapter {
 
     override fun getDisplayName(): String = message("datagrip.auth.secrets_manager")
 
-    override fun createWidget(project: Project?, creds: DatabaseCredentials, source: LocalDataSource): AuthWidget? =
-        SecretsManagerAuthWidget()
+    override fun createWidget(): AuthWidget? = SecretsManagerAuthWidget()
 
     override fun intercept(
         connection: ProtoConnection,

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/datagrip/auth/SecretsManagerAuthWidget.kt

@@ -4,7 +4,8 @@
 package software.aws.toolkits.jetbrains.datagrip.auth
 
 import com.intellij.database.dataSource.DataSourceUiUtil
-import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.database.dataSource.DatabaseConnectionConfig
+import com.intellij.database.dataSource.DatabaseConnectionPoint
 import com.intellij.ui.components.JBCheckBox
 import com.intellij.ui.components.JBLabel
 import com.intellij.ui.components.JBTextField
@@ -38,28 +39,28 @@ class SecretsManagerAuthWidget : AwsAuthWidget(userFieldEnabled = false) {
         return panel
     }
 
-    override fun save(dataSource: LocalDataSource, copyCredentials: Boolean) {
-        super.save(dataSource, copyCredentials)
+    override fun save(config: DatabaseConnectionConfig, copyCredentials: Boolean) {
+        super.save(config, copyCredentials)
 
         DataSourceUiUtil.putOrRemove(
-            dataSource.additionalProperties,
+            config.additionalProperties,
             SECRET_ID_PROPERTY,
             secretIdSelector.text.nullize()
         )
 
         DataSourceUiUtil.putOrRemove(
-            dataSource.additionalProperties,
+            config.additionalProperties,
             GET_URL_FROM_SECRET,
             urlFromSecret.isSelected.toString()
         )
     }
 
-    override fun reset(dataSource: LocalDataSource, resetCredentials: Boolean) {
-        super.reset(dataSource, resetCredentials)
-        dataSource.additionalProperties[SECRET_ID_PROPERTY]?.nullize()?.let {
+    override fun reset(config: DatabaseConnectionPoint, resetCredentials: Boolean) {
+        super.reset(config, resetCredentials)
+        config.additionalProperties[SECRET_ID_PROPERTY]?.nullize()?.let {
             secretIdSelector.text = it
         }
-        dataSource.additionalProperties[GET_URL_FROM_SECRET]?.nullize()?.let {
+        config.additionalProperties[GET_URL_FROM_SECRET]?.nullize()?.let {
             urlFromSecret.isSelected = it.toBoolean()
         }
     }

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/datagrip/auth/compatability/DatabaseAuthProviderCompatabilityAdapter.kt

@@ -3,15 +3,34 @@
 
 package software.aws.toolkits.jetbrains.datagrip.auth.compatability
 
+import com.intellij.database.access.DatabaseCredentials
 import com.intellij.database.dataSource.DatabaseAuthProvider
+import com.intellij.database.dataSource.DatabaseAuthProvider.ApplicabilityLevel
+import com.intellij.database.dataSource.DatabaseConnectionConfig
 import com.intellij.database.dataSource.DatabaseConnectionInterceptor
+import com.intellij.database.dataSource.DatabaseConnectionPoint
 import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.openapi.project.Project
 
+@Suppress("UnstableApiUsage")
 interface DatabaseAuthProviderCompatabilityAdapter : DatabaseAuthProvider {
-    override fun isApplicable(dataSource: LocalDataSource, level: DatabaseAuthProvider.ApplicabilityLevel): Boolean =
-        isApplicable(dataSource)
+    override fun getApplicability(
+        point: DatabaseConnectionPoint,
+        level: ApplicabilityLevel,
+    ): ApplicabilityLevel.Result {
+        if (!isApplicable(point.dataSource)) return ApplicabilityLevel.Result.NOT_APPLICABLE
+        return super.getApplicability(point, level)
+    }
+
+    override fun createWidget(
+        project: Project?,
+        credentials: DatabaseCredentials,
+        config: DatabaseConnectionConfig,
+    ): DatabaseAuthProvider.AuthWidget? =
+        createWidget()
 
     fun isApplicable(dataSource: LocalDataSource): Boolean
+    fun createWidget(): DatabaseAuthProvider.AuthWidget?
 }
 
 fun DatabaseConnectionInterceptor.ProtoConnection.project() = project

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/rds/auth/IamAuth.kt

@@ -4,7 +4,6 @@
 package software.aws.toolkits.jetbrains.services.rds.auth
 
 import com.intellij.credentialStore.Credentials
-import com.intellij.database.access.DatabaseCredentials
 import com.intellij.database.connection.throwable.KnownDatabaseException
 import com.intellij.database.connection.throwable.info.ErrorInfo
 import com.intellij.database.connection.throwable.info.SimpleErrorInfo
@@ -58,7 +57,7 @@ class IamAuth : DatabaseAuthProviderCompatabilityAdapter {
 
     override fun isApplicable(dataSource: LocalDataSource): Boolean = iamIsApplicable(dataSource)
 
-    override fun createWidget(project: Project?, credentials: DatabaseCredentials, dataSource: LocalDataSource): AuthWidget? = IamAuthWidget()
+    override fun createWidget(): AuthWidget? = IamAuthWidget()
 
     inner class SsoNoTokenFix(val project: Project, val connection: ProtoConnection) : ErrorInfo.Fix {
 

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/rds/auth/IamAuthWidget.kt

@@ -4,7 +4,8 @@
 package software.aws.toolkits.jetbrains.services.rds.auth
 
 import com.intellij.database.dataSource.DataSourceUiUtil
-import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.database.dataSource.DatabaseConnectionConfig
+import com.intellij.database.dataSource.DatabaseConnectionPoint
 import com.intellij.database.dataSource.url.template.ParametersHolder
 import com.intellij.ui.components.JBLabel
 import com.intellij.ui.components.JBTextField
@@ -37,8 +38,8 @@ class IamAuthWidget : AwsAuthWidget() {
         return panel
     }
 
-    override fun save(dataSource: LocalDataSource, copyCredentials: Boolean) {
-        super.save(dataSource, copyCredentials)
+    override fun save(config: DatabaseConnectionConfig, copyCredentials: Boolean) {
+        super.save(config, copyCredentials)
 
         // If the user has not specified a signing host/port we will try to use the URL in the connection
         val host = if (rdsSigningHostField.text.isNullOrBlank()) {
@@ -47,7 +48,7 @@ class IamAuthWidget : AwsAuthWidget() {
             rdsSigningHostField.text
         }
         DataSourceUiUtil.putOrRemove(
-            dataSource.additionalProperties,
+            config.additionalProperties,
             RDS_SIGNING_HOST_PROPERTY,
             host
         )
@@ -58,16 +59,16 @@ class IamAuthWidget : AwsAuthWidget() {
             rdsSigningPortField.text
         }
         DataSourceUiUtil.putOrRemove(
-            dataSource.additionalProperties,
+            config.additionalProperties,
             RDS_SIGNING_PORT_PROPERTY,
             port
         )
     }
 
-    override fun reset(dataSource: LocalDataSource, resetCredentials: Boolean) {
-        super.reset(dataSource, resetCredentials)
-        rdsSigningHostField.text = dataSource.additionalProperties[RDS_SIGNING_HOST_PROPERTY]
-        rdsSigningPortField.text = dataSource.additionalProperties[RDS_SIGNING_PORT_PROPERTY]
+    override fun reset(config: DatabaseConnectionPoint, resetCredentials: Boolean) {
+        super.reset(config, resetCredentials)
+        rdsSigningHostField.text = config.additionalProperties[RDS_SIGNING_HOST_PROPERTY]
+        rdsSigningPortField.text = config.additionalProperties[RDS_SIGNING_PORT_PROPERTY]
     }
 
     override fun updateFromUrl(holder: ParametersHolder) {

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/redshift/auth/IamAuth.kt

@@ -4,12 +4,10 @@
 package software.aws.toolkits.jetbrains.services.redshift.auth
 
 import com.intellij.credentialStore.Credentials
-import com.intellij.database.access.DatabaseCredentials
 import com.intellij.database.dataSource.DatabaseAuthProvider.AuthWidget
 import com.intellij.database.dataSource.DatabaseConnectionInterceptor.ProtoConnection
 import com.intellij.database.dataSource.DatabaseCredentialsAuthProvider
 import com.intellij.database.dataSource.LocalDataSource
-import com.intellij.openapi.project.Project
 import kotlinx.coroutines.future.future
 import software.amazon.awssdk.services.redshift.RedshiftClient
 import software.aws.toolkits.core.ConnectionSettings
@@ -37,7 +35,7 @@ class IamAuth : DatabaseAuthProviderCompatabilityAdapter {
     override fun isApplicable(dataSource: LocalDataSource): Boolean = dataSource.dbms.isRedshift
     override fun getDisplayName(): String = message("redshift.auth.aws")
 
-    override fun createWidget(project: Project?, creds: DatabaseCredentials, source: LocalDataSource): AuthWidget? = IamAuthWidget()
+    override fun createWidget(): AuthWidget? = IamAuthWidget()
 
     override fun intercept(connection: ProtoConnection, silent: Boolean): CompletionStage<ProtoConnection>? {
         LOG.info { "Intercepting db connection [$connection]" }

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/services/redshift/auth/IamAuthWidget.kt

@@ -4,7 +4,8 @@
 package software.aws.toolkits.jetbrains.services.redshift.auth
 
 import com.intellij.database.dataSource.DataSourceUiUtil
-import com.intellij.database.dataSource.LocalDataSource
+import com.intellij.database.dataSource.DatabaseConnectionConfig
+import com.intellij.database.dataSource.DatabaseConnectionPoint
 import com.intellij.database.dataSource.url.template.ParametersHolder
 import com.intellij.database.dataSource.url.template.UrlEditorModel
 import com.intellij.ui.components.JBLabel
@@ -34,19 +35,19 @@ class IamAuthWidget : AwsAuthWidget() {
         return panel
     }
 
-    override fun save(dataSource: LocalDataSource, copyCredentials: Boolean) {
-        super.save(dataSource, copyCredentials)
+    override fun save(config: DatabaseConnectionConfig, copyCredentials: Boolean) {
+        super.save(config, copyCredentials)
 
         DataSourceUiUtil.putOrRemove(
-            dataSource.additionalProperties,
+            config.additionalProperties,
             CLUSTER_ID_PROPERTY,
             clusterIdSelector.text
         )
     }
 
-    override fun reset(dataSource: LocalDataSource, resetCredentials: Boolean) {
-        super.reset(dataSource, resetCredentials)
-        clusterIdSelector.text = dataSource.additionalProperties[CLUSTER_ID_PROPERTY]
+    override fun reset(config: DatabaseConnectionPoint, resetCredentials: Boolean) {
+        super.reset(config, resetCredentials)
+        clusterIdSelector.text = config.additionalProperties[CLUSTER_ID_PROPERTY]
     }
 
     override fun updateFromUrl(holder: ParametersHolder) {

plugins/toolkit/jetbrains-ultimate/src/software/aws/toolkits/jetbrains/ui/AwsAuthWidget.kt

@@ -1,19 +1,124 @@
-// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 package software.aws.toolkits.jetbrains.ui
 
+import com.intellij.database.dataSource.DataSourceUiUtil
 import com.intellij.database.dataSource.DatabaseConnectionConfig
 import com.intellij.database.dataSource.DatabaseConnectionPoint
+import com.intellij.database.dataSource.DatabaseCredentialsAuthProviderUi
+import com.intellij.database.dataSource.url.template.ParametersHolder
+import com.intellij.database.dataSource.url.template.UrlEditorModel
+import com.intellij.ui.components.JBLabel
+import com.intellij.uiDesigner.core.GridLayoutManager
+import com.intellij.util.text.nullize
+import org.jetbrains.annotations.TestOnly
+import software.aws.toolkits.jetbrains.UiConstraints
+import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
+import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
+import software.aws.toolkits.jetbrains.datagrip.CREDENTIAL_ID_PROPERTY
+import software.aws.toolkits.jetbrains.datagrip.REGION_ID_PROPERTY
+import software.aws.toolkits.jetbrains.utils.ui.selected
+import software.aws.toolkits.resources.message
+import javax.swing.JPanel
+
+abstract class AwsAuthWidget(private val userFieldEnabled: Boolean = true) : DatabaseCredentialsAuthProviderUi.UserWidget() {
+    private val credentialSelector = CredentialProviderSelector()
+    private val regionSelector = RegionSelector()
+
+    // These are defined in DataGrip and used in the url holder which is called in `updateFromUrl`
+    protected val hostParameter = "host"
+    protected val portParameter = "port"
+
+    abstract fun getRegionFromUrl(url: String?): String?
+    abstract val serviceId: String
+
+    open val rowCount: Int = 3
+    open val columnCount: Int = 6
+
+    override fun createPanel(): JPanel {
+        val panel = JPanel(GridLayoutManager(rowCount, columnCount))
+        addUserField(panel, 0)
+
+        // Disable the user field if we treat it as immutable
+        myUserField.isEnabled = userFieldEnabled
+
+        val credsLabel = JBLabel(message("aws_connection.credentials.label"))
+        val regionLabel = JBLabel(message("aws_connection.region.label"))
+        panel.add(credsLabel, UiConstraints.createLabelConstraints(1, 0, credsLabel.preferredSize.getWidth()))
+        panel.add(credentialSelector, UiConstraints.createSimpleConstraints(1, 1, 3))
+        panel.add(regionLabel, UiConstraints.createLabelConstraints(2, 0, regionLabel.preferredSize.getWidth()))
+        panel.add(regionSelector, UiConstraints.createSimpleConstraints(2, 1, 3))
+
+        return panel
+    }
 
-abstract class AwsAuthWidget(userFieldEnabled: Boolean = true) : @Suppress("DEPRECATION") AwsAuthWidgetBase(userFieldEnabled) {
     override fun save(config: DatabaseConnectionConfig, copyCredentials: Boolean) {
         super.save(config, copyCredentials)
-        save(config.dataSource, copyCredentials)
+
+        DataSourceUiUtil.putOrRemove(
+            config.additionalProperties,
+            CREDENTIAL_ID_PROPERTY,
+            credentialSelector.getSelectedCredentialsProvider()
+        )
+        DataSourceUiUtil.putOrRemove(
+            config.additionalProperties,
+            REGION_ID_PROPERTY,
+            regionSelector.selectedRegion?.id
+        )
     }
 
-    override fun reset(point: DatabaseConnectionPoint, resetCredentials: Boolean) {
-        super.reset(point, resetCredentials)
-        reset(point.dataSource, resetCredentials)
+    override fun reset(config: DatabaseConnectionPoint, resetCredentials: Boolean) {
+        super.reset(config, resetCredentials)
+
+        val regionProvider = AwsRegionProvider.getInstance()
+        val allRegions = regionProvider.allRegionsForService(serviceId)
+        regionSelector.setRegions(allRegions.values.toMutableList())
+        val regionId = config.additionalProperties[REGION_ID_PROPERTY]?.nullize()
+        regionId?.let {
+            allRegions[regionId]?.let { region ->
+                regionSelector.selectedRegion = region
+            }
+        }
+
+        val credentialManager = CredentialManager.getInstance()
+        credentialSelector.setCredentialsProviders(credentialManager.getCredentialIdentifiers())
+        val credentialId = config.additionalProperties[CREDENTIAL_ID_PROPERTY]?.nullize()
+        if (credentialId != null) {
+            val credentialIdentifierById = credentialManager.getCredentialIdentifierById(credentialId)
+            if (credentialIdentifierById != null) {
+                credentialSelector.setSelectedCredentialsProvider(credentialIdentifierById)
+            } else {
+                credentialSelector.setSelectedInvalidCredentialsProvider(credentialId)
+            }
+        } else {
+            credentialSelector.model.selectedItem = null
+        }
+    }
+
+    override fun isPasswordChanged(): Boolean = false
+
+    override fun onChanged(r: Runnable) {
+        // Tries to set username so if we don't have one, don't set
+        if (userFieldEnabled) {
+            super.onChanged(r)
+        }
     }
+
+    override fun updateFromUrl(holder: ParametersHolder) {
+        // Try to get region from url and set the region box on a best effort basis
+        val url = (holder as? UrlEditorModel)?.url
+        val regionId = getRegionFromUrl(url)
+        val region = AwsRegionProvider.getInstance().allRegions()[regionId]
+        region?.let {
+            regionSelector.selectedRegion = it
+        }
+        super.updateFromUrl(holder)
+    }
+
+    @TestOnly
+    internal fun getSelectedCredential() = credentialSelector.getSelectedCredentialsProvider()
+
+    @TestOnly
+    internal fun getSelectedRegion() = regionSelector.selected()
 }