Set data source properties for AWS secret fetching (#3485)

* Set data source properties for AWS secret fetching

* Update misk-jdbc/src/main/kotlin/misk/jdbc/DataSourceConfig.kt

Co-authored-by: Andrew (Paradi) Alexander <[email protected]>

---------

Co-authored-by: Andrew (Paradi) Alexander <[email protected]>

Author: katukota

misk-jdbc/api/misk-jdbc.api

@@ -303,6 +303,7 @@ public final class misk/jdbc/DataSourceConfig {
 	public final fun getConnection_idle_timeout ()Ljava/time/Duration;
 	public final fun getConnection_max_lifetime ()Ljava/time/Duration;
 	public final fun getConnection_timeout ()Ljava/time/Duration;
+	public final fun getDataSourceProperties ()Ljava/util/Properties;
 	public final fun getDatabase ()Ljava/lang/String;
 	public final fun getDeclarative_schema_config ()Lmisk/jdbc/DeclarativeSchemaConfig;
 	public final fun getDriverClassName ()Ljava/lang/String;

misk-jdbc/src/main/kotlin/misk/jdbc/DataSourceConfig.kt

@@ -5,6 +5,7 @@ import misk.config.Config
 import misk.containers.ContainerUtil
 import wisp.deployment.Deployment
 import java.time.Duration
+import java.util.Properties
 
 /** Defines a type of datasource */
 enum class DataSourceType(
@@ -156,6 +157,30 @@ data class DataSourceConfig @JvmOverloads constructor(
     }
   }
 
+  fun getDataSourceProperties() : Properties {
+    // https://github.com/brettwooldridge/HikariCP/wiki/MySQL-Configuration
+    val properties = Properties()
+
+    properties["cachePrepStmts"] = "true"
+    properties["prepStmtCacheSize"] = "250"
+    properties["prepStmtCacheSqlLimit"] = "2048"
+    if (type == DataSourceType.MYSQL || type == DataSourceType.VITESS_MYSQL || type == DataSourceType.TIDB) {
+      properties["useServerPrepStmts"] = "true"
+    }
+    if (mysql_use_aws_secret_for_credentials) {
+      properties["user"] = mysql_aws_secret_name
+    }
+    properties["useLocalSessionState"] = "true"
+    properties["rewriteBatchedStatements"] = "true"
+    properties["cacheResultSetMetadata"] = "true"
+    properties["cacheServerConfiguration"] = "true"
+    properties["elideSetAutoCommits"] = "true"
+    properties["maintainTimeStats"] = "false"
+    properties["characterEncoding"] = "UTF-8"
+
+    return properties
+  }
+
   fun withDefaults(): DataSourceConfig {
     val server_hostname = ContainerUtil.dockerTargetOrLocalIp()
     return when (type) {
@@ -294,8 +319,6 @@ data class DataSourceConfig @JvmOverloads constructor(
         }
 
         if(mysql_use_aws_secret_for_credentials) {
-          val region = "us-east-1"
-          queryParams += "&secretId=$mysql_aws_secret_name&region=$region"
           "jdbc-secretsmanager:mysql://${config.host}:${config.port}/${config.database}$queryParams"
         }
         else {

misk-jdbc/src/main/kotlin/misk/jdbc/DataSourceService.kt

@@ -105,21 +105,7 @@ class DataSourceService @JvmOverloads constructor(
         else -> {}
       }
 
-
-      // https://github.com/brettwooldridge/HikariCP/wiki/MySQL-Configuration
-      hikariConfig.dataSourceProperties["cachePrepStmts"] = "true"
-      hikariConfig.dataSourceProperties["prepStmtCacheSize"] = "250"
-      hikariConfig.dataSourceProperties["prepStmtCacheSqlLimit"] = "2048"
-      if (config.type == DataSourceType.MYSQL || config.type == DataSourceType.VITESS_MYSQL || config.type == DataSourceType.TIDB) {
-        hikariConfig.dataSourceProperties["useServerPrepStmts"] = "true"
-      }
-      hikariConfig.dataSourceProperties["useLocalSessionState"] = "true"
-      hikariConfig.dataSourceProperties["rewriteBatchedStatements"] = "true"
-      hikariConfig.dataSourceProperties["cacheResultSetMetadata"] = "true"
-      hikariConfig.dataSourceProperties["cacheServerConfiguration"] = "true"
-      hikariConfig.dataSourceProperties["elideSetAutoCommits"] = "true"
-      hikariConfig.dataSourceProperties["maintainTimeStats"] = "false"
-      hikariConfig.dataSourceProperties["characterEncoding"] = "UTF-8"
+      hikariConfig.dataSourceProperties = config.getDataSourceProperties()
     }
 
     // TODO(sahilm): The same mitigation _might_ be applicable to the DataSourceTypes VITESS_MYSQL and TIDB

misk-jdbc/src/main/kotlin/misk/jdbc/PingDatabaseService.kt

@@ -7,7 +7,6 @@ import misk.backoff.retry
 import wisp.deployment.Deployment
 import misk.logging.getLogger
 import java.time.Duration
-import java.util.*
 import jakarta.inject.Inject
 import jakarta.inject.Singleton
 import misk.backoff.RetryConfig
@@ -65,7 +64,9 @@ class PingDatabaseService @Inject constructor(
 
   private fun createDataSource(jdbcUrl: String): DriverDataSource {
     return DriverDataSource(
-            jdbcUrl, config.getDriverClassName(), Properties(), config.username, config.password)
+      jdbcUrl, config.getDriverClassName(), config.getDataSourceProperties(),
+      config.username, config.password
+    )
   }
 
   /** Kotlin thinks getConnection() is a val but it's really a function. */

misk-jdbc/src/test/kotlin/misk/jdbc/DataSourceConfigTest.kt

@@ -300,8 +300,7 @@ class DataSourceConfigTest {
     assertEquals(
       "jdbc-secretsmanager:mysql://127.0.0.1:3306/?useLegacyDatetimeCode=false&" +
           "createDatabaseIfNotExist=true&connectTimeout=10000&socketTimeout=60000&" +
-          "sslMode=PREFERRED&enabledTLSProtocols=TLSv1.2,TLSv1.3&" +
-          "secretId=secret_name&region=us-east-1",
+          "sslMode=PREFERRED&enabledTLSProtocols=TLSv1.2,TLSv1.3",
       config.buildJdbcUrl(TESTING)
     )
   }
@@ -321,7 +320,6 @@ class DataSourceConfigTest {
     // Should generate correct JDBC URL
     val jdbcUrl = config.buildJdbcUrl(TESTING)
     assertThat(jdbcUrl).startsWith("jdbc-secretsmanager:mysql://")
-    assertThat(jdbcUrl).contains("secretId=test-secret")
   }
 
   @Test