refetching of rotated credentials with fallback to config values

Author: salamonpavel

server/src/main/resources/reference.conf

@@ -1,19 +1,18 @@
 {
   postgres {
-    # The JDBC driver class
-    dataSourceClass=org.postgresql.Driver
+    dataSourceClass="za.co.absa.atum.server.api.database.PostgresDataSourceWithPasswordFromSecretsManager"
     # host.docker.internal for local run against db running in docker on its host machine; localhost otherwise for testing and for the gh pipeline
     serverName=localhost
     portNumber=5432
     databaseName=atum_db
     user=atum_user
     password=changeme
+    passwordSecretId="bdtools-atum-service-dev/atum_service_user_password"
     # maximum number of connections that HikariCP will keep in the pool, including both idle and in-use connections
     maxPoolSize=10
   }
   aws {
     region = "af-south-1"
-    dbPasswordSecretName = "serviceUserSecretKey"
   }
   ssl {
     enabled=false

server/src/main/scala/za/co/absa/atum/server/Main.scala

@@ -53,8 +53,6 @@ import za.co.absa.atum.server.api.v2.repository.{
 import za.co.absa.atum.server.api.database.flows.functions.{GetFlowCheckpoints, GetFlowPartitionings}
 import za.co.absa.atum.server.api.database.runs.functions._
 import za.co.absa.atum.server.api.database.{PostgresDatabaseProvider, TransactorProvider}
-
-import za.co.absa.atum.server.aws.AwsSecretsProviderImpl
 import za.co.absa.atum.server.config.JvmMonitoringConfig
 import zio._
 import zio.config.typesafe.TypesafeConfigProvider
@@ -111,8 +109,6 @@ object Main extends ZIOAppDefault {
           GetCheckpointProperties.layer,
           PostgresDatabaseProvider.layer,
           TransactorProvider.layer,
-          // aws
-          AwsSecretsProviderImpl.layer,
           // scope
           zio.Scope.default,
           // prometheus

server/src/main/scala/za/co/absa/atum/server/api/database/AWSSDKs.scala

@@ -0,0 +1,23 @@
+package za.co.absa.atum.server.api.database
+
+import com.typesafe.config.ConfigFactory
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider
+import software.amazon.awssdk.regions.Region
+import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient
+
+object AWSSDKs {
+
+  private val config = ConfigFactory
+    .load()
+    .getConfig(s"aws")
+
+  private val defaultRegion = config.getString("region")
+
+  private val credentialsProvider = DefaultCredentialsProvider.create
+
+  val secretsManagerSyncClient: SecretsManagerClient = SecretsManagerClient.builder
+    .region(Region.of(defaultRegion))
+    .credentialsProvider(credentialsProvider)
+    .build
+
+}

server/src/main/scala/za/co/absa/atum/server/api/database/PostgresDataSourceWithPasswordFromSecretsManager.scala

@@ -0,0 +1,106 @@
+package za.co.absa.atum.server.api.database
+
+import com.typesafe.config.ConfigFactory
+import org.postgresql.ds.PGSimpleDataSource
+import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient
+import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest
+
+import java.sql.{Connection, SQLException}
+import scala.util.{Failure, Try}
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+
+/**
+ *  `PGSimpleDataSource` but with password being fetched from AWS Secrets Manager.
+ *
+ *  Refreshes the password when `PGSimpleDataSource.getConnection` fails,
+ *  thus it works with secrets with rotation enabled.
+ *
+ *  Expects the same set of properties as `PGSimpleDataSource` +
+ *  `passwordSecretId` - ID of the secret containing password.
+ */
+class PostgresDataSourceWithPasswordFromSecretsManager extends PGSimpleDataSource {
+
+  protected val logger: Logger = LoggerFactory.getLogger(this.getClass())
+
+  private var password: String = _
+  private var passwordSecretId: String = _
+
+  override def getConnection(): Connection = {
+    if (Option(password).isEmpty) {
+      val pw = getPasswordFromSecretsManagerOrConfig()
+      setInternalPassword(pw)
+    }
+
+    val connectionTry = Try(baseGetConnection(user, password)).recoverWith { case _: SQLException =>
+      logger.info("Failed to create Postgres connection, attempting to refresh the password and try again...")
+      for {
+        passwordFromSecretsManager <- getPasswordFromSecretsManager()
+        connection <- Try(baseGetConnection(user, passwordFromSecretsManager)).recoverWith { case e =>
+          logger.error("Failed to create Postgres connection even after password refresh")
+          Failure(e)
+        }
+      } yield {
+        setInternalPassword(passwordFromSecretsManager)
+        connection
+      }
+    }
+
+    connectionTry.get
+  }
+
+  // getter and setter for passwordSecretId are needed as this class is usually constructed by reflection
+  def getPasswordSecretId(): String = passwordSecretId
+
+  def setPasswordSecretId(passwordSecretId: String): Unit = {
+    this.passwordSecretId = passwordSecretId
+  }
+
+  override def setProperty(name: String, value: String): Unit = name match {
+    case "passwordSecretId" => setPasswordSecretId(value)
+    case _ => baseSetProperty(name, value)
+  }
+
+  // the following protected defs are for easier unit tests
+  protected def baseSetProperty(name: String, value: String): Unit = super.setProperty(name, value)
+  protected def baseGetConnection(username: String, password: String): Connection =
+    super.getConnection(username, password)
+  protected def user: String = this.getUser
+  protected def secretsManagerClient: SecretsManagerClient = AWSSDKs.secretsManagerSyncClient
+
+  private[database] def setInternalPassword(password: String): Unit = {
+    this.password = password
+  }
+
+  private def getPasswordFromSecretsManager(): Try[String] = {
+    val secretID = getPasswordSecretId()
+
+    val secretValueTry = Try {
+      logger.info(s"Fetching password for Postgres from Secrets Manager (secret id: $secretID)")
+      val response = secretsManagerClient.getSecretValue(
+        GetSecretValueRequest.builder
+          .secretId(secretID)
+          .build
+      )
+      logger.info("Successfully fetched password for Postgres from Secrets Manager")
+      response.secretString
+    }
+
+    secretValueTry.recoverWith { case e =>
+      logger.error(s"Failed to fetch password for Postgres from Secrets Manager (secret id: $secretID)")
+      Failure(e)
+    }
+  }
+
+  private def getPasswordFromSecretsManagerOrConfig(): String = {
+    getPasswordFromSecretsManager().getOrElse {
+      logger.error(
+        s"Failed to fetch password from Secrets Manager (secret id: ${getPasswordSecretId()}). " +
+          s"Falling back to config value."
+      )
+      val configPassword = ConfigFactory.load().getConfig("postgres").getString("password")
+      configPassword
+    }
+  }
+
+}

server/src/main/scala/za/co/absa/atum/server/api/database/TransactorProvider.scala

@@ -18,38 +18,30 @@ package za.co.absa.atum.server.api.database
 
 import com.zaxxer.hikari.HikariConfig
 import doobie.hikari.HikariTransactor
-import za.co.absa.atum.server.aws.AwsSecretsProvider
-import za.co.absa.atum.server.config.{AwsConfig, PostgresConfig}
+import za.co.absa.atum.server.config.PostgresConfig
 import zio.Runtime.defaultBlockingExecutor
 import zio._
 import zio.interop.catz._
 
 object TransactorProvider {
 
-  val layer: ZLayer[Any with Scope with AwsSecretsProvider, Throwable, HikariTransactor[Task]] = ZLayer {
+  val layer: ZLayer[Any with Scope, Throwable, HikariTransactor[Task]] = ZLayer {
     for {
       postgresConfig <- ZIO.config[PostgresConfig](PostgresConfig.config)
-      awsConfig <- ZIO.config[AwsConfig](AwsConfig.config)
-
-      awsSecretsProvider <- ZIO.service[AwsSecretsProvider]
-      password <- awsSecretsProvider
-        .getSecretValue(awsConfig.dbPasswordSecretName)
-        // fallback to password property's value from postgres section of reference.conf; useful for local testing
-        .orElse {
-          ZIO
-            .logError("Credentials were not retrieved from AWS, falling back to config value.")
-            .as(postgresConfig.password)
-        }
 
       hikariConfig = {
+        val dataSourceProperties = new java.util.Properties()
+        dataSourceProperties.setProperty("serverName", postgresConfig.serverName)
+        dataSourceProperties.setProperty("portNumber", postgresConfig.portNumber.toString)
+        dataSourceProperties.setProperty("databaseName", postgresConfig.databaseName)
+        dataSourceProperties.setProperty("user", postgresConfig.user)
+        dataSourceProperties.setProperty("passwordSecretId", postgresConfig.passwordSecretId)
+
         val config = new HikariConfig()
-        config.setDriverClassName(postgresConfig.dataSourceClass)
-        config.setJdbcUrl(
-          s"jdbc:postgresql://${postgresConfig.serverName}:${postgresConfig.portNumber}/${postgresConfig.databaseName}"
-        )
-        config.setUsername(postgresConfig.user)
-        config.setPassword(password)
+        config.setDataSourceClassName(postgresConfig.dataSourceClass)
+        config.setDataSourceProperties(dataSourceProperties)
         config.setMaximumPoolSize(postgresConfig.maxPoolSize)
+        config.setPoolName("DoobiePostgresHikariPool")
         config
       }
 

server/src/main/scala/za/co/absa/atum/server/aws/AwsSecretsProvider.scala

@@ -26,7 +26,8 @@ case class PostgresConfig(
   databaseName: String,
   user: String,
   password: String,
-  maxPoolSize: Int
+  maxPoolSize: Int,
+  passwordSecretId: String
 )
 
 object PostgresConfig {

server/src/main/scala/za/co/absa/atum/server/aws/AwsSecretsProviderImpl.scala

@@ -1,19 +1,18 @@
 {
   postgres {
-      # The JDBC driver class
-      dataSourceClass=org.postgresql.Driver
+      dataSourceClass="za.co.absa.atum.server.api.database.PostgresDataSourceWithPasswordFromSecretsManager"
       serverName=localhost
       portNumber=5432
       databaseName=atum_db
       # tests have to be run with atum_owner so we can execute not only plpgsql functions granted to atum_user
       user=atum_owner
       password=changeme
+      passwordSecretId="bdtools-atum-service-dev/atum_service_user_password"
       # maximum number of connections that HikariCP will keep in the pool, including both idle and in-use connections
       maxPoolSize=10
    }
   aws {
     region = "af-south-1"
-    dbPasswordSecretName = "serviceUserSecretKey"
   }
   ssl {
     enabled=false