Fix IllegalStateException when using connected execution experiement is enabled (#3007)

Author: rli

jetbrains-core/src/software/aws/toolkits/jetbrains/core/execution/AwsConnectionExtension.kt

@@ -16,6 +16,7 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.tryOrNull
 import software.aws.toolkits.jetbrains.core.credentials.AwsConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.CredentialManager
+import software.aws.toolkits.jetbrains.core.execution.AwsCredentialInjectionOptions.Companion.DEFAULT_OPTIONS
 import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.AwsTelemetry
@@ -24,26 +25,11 @@ import software.aws.toolkits.telemetry.Result.Succeeded
 
 class AwsConnectionRunConfigurationExtension<T : RunConfigurationBase<*>> {
     fun addEnvironmentVariables(configuration: T, environment: MutableMap<String, String>, runtimeString: () -> String? = { null }) {
-        val regionProvider = AwsRegionProvider.getInstance()
-        val credentialManager = CredentialManager.getInstance()
-        val credentialConfiguration = configuration.getCopyableUserData(AWS_CONNECTION_RUN_CONFIGURATION_KEY) ?: return
-
         try {
-            val connection = if (credentialConfiguration.useCurrentConnection) {
-                AwsConnectionManager.getInstance(configuration.project).connectionSettings() ?: throw RuntimeException(message("configure.toolkit"))
-            } else {
-                val region = credentialConfiguration.region?.let {
-                    regionProvider.allRegions()[it]
-                } ?: throw IllegalStateException(message("configure.validate.no_region_specified"))
-
-                val credentialProviderId = credentialConfiguration.credential ?: throw IllegalStateException(message("aws.notification.credentials_missing"))
-
-                val credentialProvider = credentialManager.getCredentialIdentifierById(credentialProviderId)?.let {
-                    credentialManager.getAwsCredentialProvider(it, region)
-                } ?: throw RuntimeException(message("aws.notification.credentials_missing"))
+            val credentialConfiguration = credentialConfiguration(configuration) ?: return
+            if (credentialConfiguration == DEFAULT_OPTIONS) return
 
-                ConnectionSettings(credentialProvider, region)
-            }
+            val connection = getConnection(configuration, credentialConfiguration)
             connection.region.mergeWithExistingEnvironmentVariables(environment)
             connection.credentials.resolveCredentials().mergeWithExistingEnvironmentVariables(environment)
             AwsTelemetry.injectCredentials(configuration.project, result = Succeeded, runtimeString = tryOrNull { runtimeString() })
@@ -53,6 +39,13 @@ class AwsConnectionRunConfigurationExtension<T : RunConfigurationBase<*>> {
         }
     }
 
+    fun validateConfiguration(runConfiguration: T, isExecution: Boolean) {
+        val credentialConfiguration = runConfiguration.getCopyableUserData(AWS_CONNECTION_RUN_CONFIGURATION_KEY) ?: return
+        if (credentialConfiguration == DEFAULT_OPTIONS) return
+
+        getConnection(runConfiguration, credentialConfiguration)
+    }
+
     fun readExternal(runConfiguration: T, element: Element) {
         runConfiguration.putCopyableUserData(
             AWS_CONNECTION_RUN_CONFIGURATION_KEY,
@@ -69,6 +62,29 @@ class AwsConnectionRunConfigurationExtension<T : RunConfigurationBase<*>> {
         }
     }
 
+    private fun credentialConfiguration(configuration: T) = configuration.getCopyableUserData(AWS_CONNECTION_RUN_CONFIGURATION_KEY)
+
+    private fun getConnection(configuration: T, credentialConfiguration: AwsCredentialInjectionOptions): ConnectionSettings {
+        val regionProvider = AwsRegionProvider.getInstance()
+        val credentialManager = CredentialManager.getInstance()
+
+        return if (credentialConfiguration.useCurrentConnection) {
+            AwsConnectionManager.getInstance(configuration.project).connectionSettings() ?: throw RuntimeException(message("configure.toolkit"))
+        } else {
+            val region = credentialConfiguration.region?.let {
+                regionProvider.allRegions()[it]
+            } ?: throw IllegalStateException(message("configure.validate.no_region_specified"))
+
+            val credentialProviderId = credentialConfiguration.credential ?: throw IllegalStateException(message("aws.notification.credentials_missing"))
+
+            val credentialProvider = credentialManager.getCredentialIdentifierById(credentialProviderId)?.let {
+                credentialManager.getAwsCredentialProvider(it, region)
+            } ?: throw RuntimeException(message("aws.notification.credentials_missing"))
+
+            ConnectionSettings(credentialProvider, region)
+        }
+    }
+
     private companion object {
         val LOG = getLogger<AwsConnectionRunConfigurationExtension<*>>()
     }
@@ -88,12 +104,13 @@ val AWS_CONNECTION_RUN_CONFIGURATION_KEY =
         "aws.toolkit.runConfigurationConnection"
     )
 
-class AwsCredentialInjectionOptions {
-    var useCurrentConnection: Boolean = false
-    var region: String? = null
+data class AwsCredentialInjectionOptions(
+    var useCurrentConnection: Boolean = false,
+    var region: String? = null,
     var credential: String? = null
-
+) {
     companion object {
         operator fun invoke(block: AwsCredentialInjectionOptions.() -> Unit) = AwsCredentialInjectionOptions().apply(block)
+        val DEFAULT_OPTIONS by lazy { AwsCredentialInjectionOptions() }
     }
 }

jetbrains-core/src/software/aws/toolkits/jetbrains/core/execution/JavaAwsConnectionExtension.kt

@@ -41,6 +41,10 @@ class JavaAwsConnectionExtension : RunConfigurationExtension() {
 
     override fun <T : RunConfigurationBase<*>?> createEditor(configuration: T): SettingsEditor<T>? = connectionSettingsEditor(configuration)
 
+    override fun validateConfiguration(configuration: RunConfigurationBase<*>, isExecution: Boolean) {
+        delegate.validateConfiguration(configuration, isExecution)
+    }
+
     override fun readExternal(runConfiguration: RunConfigurationBase<*>, element: Element) = delegate.readExternal(runConfiguration, element)
 
     override fun writeExternal(runConfiguration: RunConfigurationBase<*>, element: Element) = delegate.writeExternal(runConfiguration, element)

jetbrains-core/src/software/aws/toolkits/jetbrains/core/execution/PythonAwsConnectionExtension.kt

@@ -41,6 +41,10 @@ class PythonAwsConnectionExtension : PythonRunConfigurationExtension() {
         configuration
     )
 
+    override fun validateConfiguration(configuration: AbstractPythonRunConfiguration<*>, isExecution: Boolean) {
+        delegate.validateConfiguration(configuration, isExecution)
+    }
+
     private fun isEnabled() = PythonAwsConnectionExperiment.isEnabled()
 }
 

jetbrains-core/tst/software/aws/toolkits/jetbrains/core/execution/JavaAwsConnectionExtensionTest.kt

@@ -91,6 +91,16 @@ class JavaAwsConnectionExtensionTest {
         assertThat(configuration.getCopyableUserData(AWS_CONNECTION_RUN_CONFIGURATION_KEY)).isNull()
     }
 
+    @Test
+    fun `Does not throw on default options`() {
+        val runManager = RunManager.getInstance(projectRule.project)
+        val configuration = runManager.createConfiguration("test", ApplicationConfigurationType::class.java).configuration as ApplicationConfiguration
+        configuration.putCopyableUserData(AWS_CONNECTION_RUN_CONFIGURATION_KEY, AwsCredentialInjectionOptions.DEFAULT_OPTIONS)
+        val extension = JavaAwsConnectionExtension()
+        val map = mutableMapOf<String, String>()
+        extension.updateJavaParameters(configuration, mock { on { env } doAnswer { map } }, null)
+    }
+
     @Test
     fun `Inject injects environment variables`() {
         val runManager = RunManager.getInstance(projectRule.project)