Fixes issue-11740: Added support for the om service to connect to AWS RDS using IAM roles (#11913)

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Changed intial implementation accordingly. Added better flexibility for different auth prodvider impl

* ISSUE-11740: Clean up unnecessary classes

* ISSUE-11740: Clean up unnecessary properties

* ISSUE-11740: Clean up unnecessary properties

* ISSUE-11740: Clean up unnecessary properties

* ISSUE-11740: Clean up unnecessary properties

* ISSUE-11740: Clean up unnecessary properties

* ISSUE-11740: Code formatting

* ISSUE-11740: Added support for the om service to connect to AWS RDS using IAM roles

* ISSUE-11740: Moved docs to 1.2 version

---------

Co-authored-by: artiom.darie <artiom.darie@adswizz.com>
Co-authored-by: Sriharsha Chintalapani <harshach@users.noreply.github.com>

.pre-commit-config.yaml

@@ -1,9 +1,10 @@
+default_language_version:
+  python: python3
 repos:
   - repo: https://github.com/ambv/black
     rev: 22.3.0
     hooks:
       - id: black
-        language_version: python3.8
         exclude: ingestion/src/metadata/generated
   - repo: https://github.com/timothycrosley/isort
     rev: 5.12.0

conf/openmetadata.yaml

@@ -25,7 +25,7 @@ server:
       port: ${SERVER_ADMIN_PORT:-8586}
 
 # Above configuration for running http is fine for dev and testing.
-# For production setup, where UI app will hit apis through DPS it 
+# For production setup, where UI app will hit apis through DPS it
 # is strongly recommended to run https instead. Note that only
 # keyStorePath and keyStorePassword are mandatory properties. Values
 # for other properties are defaults
@@ -33,7 +33,7 @@ server:
   #applicationConnectors:
   #  - type: https
   #    port: 8585
-  #    keyStorePath: ./conf/keystore.jks 
+  #    keyStorePath: ./conf/keystore.jks
   #    keyStorePassword: changeit
   #    keyStoreType: JKS
   #    keyStoreProvider:
@@ -57,12 +57,12 @@ server:
   #    supportedCipherSuites: TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256
   #    allowRenegotiation: true
   #    endpointIdentificationAlgorithm: (none)
- 
+
   #adminConnectors:
   #  - type: https
   #    port: 8586
-  #    keyStorePath: ./conf/keystore.jks 
-  #    keyStorePassword: changeit 
+  #    keyStorePath: ./conf/keystore.jks
+  #    keyStorePassword: changeit
   #    keyStoreType: JKS
   #    keyStoreProvider:
   #    trustStorePath: /path/to/file
@@ -125,7 +125,7 @@ database:
   user: ${DB_USER:-openmetadata_user}
   password: ${DB_USER_PASSWORD:-openmetadata_password}
   # the JDBC URL; the database is called openmetadata_db
-  url: jdbc:${DB_SCHEME:-mysql}://${DB_HOST:-localhost}:${DB_PORT:-3306}/${OM_DATABASE:-openmetadata_db}?allowPublicKeyRetrieval=true&useSSL=${DB_USE_SSL:-false}&serverTimezone=UTC
+  url: jdbc:${DB_SCHEME:-mysql}://${DB_HOST:-localhost}:${DB_PORT:-3306}/${OM_DATABASE:-openmetadata_db}?${DB_PARAMS:-allowPublicKeyRetrieval=true&useSSL=false&serverTimezone=UTC}
 
 migrationConfiguration:
   flywayPath: "./bootstrap/sql/migrations/flyway"
@@ -342,7 +342,7 @@ web:
   permission-policy:
     enabled: ${WEB_CONF_PERMISSION_POLICY_ENABLED:-false}
     option: ${WEB_CONF_PERMISSION_POLICY_OPTION:-""}
-    
+
 
 changeEventConfig:
   omUri: ${OM_URI:- "http://localhost:8585"} #openmetadata in om uri for eg http://localhost:8585

openmetadata-docs/content/v0.13.3/how-to-guides/aws/index.md

@@ -0,0 +1,27 @@
+---
+title: How to enable AWS RDS IAM Auth on postgresql
+slug: /how-to-guides/aws/index.md
+---
+
+# Aws resources on Rds IAM Auth
+https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.html
+
+# Requirements
+1. AWS Rds Cluster with IAM auth enabled
+2. User on Db Cluster with iam enabled
+3. IAM policy with permission on rds connect
+4. Role with IAM policy attached
+5. IAM role attached to ec2 instance on which openmetadata is deployed or ServiceAccount/Kube2Iam role attached to pod
+  
+# How to enable ADS RDS IAM Auth on postgresql
+
+Set environment variables
+```Commandline
+  AWS_ENABLE_IAM_DATABASE_AUTHENTICATION: true
+  AWS_REGION: your_region
+  DB_PARAMS: "allowPublicKeyRetrieval=true&sslmode=require&serverTimezone=UTC"
+```
+Either through helm (if deployed in kubernetes) or as env vars
+
+# Note
+The `DB_USER_PASSWORD` is still required and cannot be empty. Set it to a random/dummy string.

openmetadata-docs/content/v1.2.x-SNAPSHOT/deployment/configuration.md

@@ -47,7 +47,6 @@ database:
   # the JDBC URL; the database is called openmetadata_db
   url: jdbc:mysql://localhost/openmetadata_db?useSSL=false&serverTimezone=UTC
 
-
 elasticsearch:
   host: localhost
   port: 9200

openmetadata-docs/content/v1.2.x-SNAPSHOT/how-to-guides/aws/index.md

@@ -0,0 +1,26 @@
+---
+title: How to enable AWS RDS IAM Auth on postgresql
+slug: /how-to-guides/aws/index.md
+---
+
+# Aws resources on Rds IAM Auth
+https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.html
+
+# Requirements
+1. AWS Rds Cluster with IAM auth enabled
+2. User on Db Cluster with iam enabled
+3. IAM policy with permission on rds connect
+4. Role with IAM policy attached
+5. IAM role attached to ec2 instance on which openmetadata is deployed or ServiceAccount/Kube2Iam role attached to pod
+  
+# How to enable ADS RDS IAM Auth on postgresql
+
+Set environment variables
+```Commandline
+  DB_USER_PASSWORD: "dummy"
+  DB_PARAMS: "awsRegion=eu-west-1&allowPublicKeyRetrieval=true&sslmode=require&serverTimezone=UTC"
+```
+Either through helm (if deployed in kubernetes) or as env vars
+
+# Note
+The `DB_USER_PASSWORD` is still required and cannot be empty. Set it to a random/dummy string.

openmetadata-service/pom.xml

@@ -247,6 +247,16 @@
       <artifactId>ssm</artifactId>
       <version>${awssdk.version}</version>
     </dependency>
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>rds</artifactId>
+      <version>${awssdk.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>sts</artifactId>
+      <version>${awssdk.version}</version>
+    </dependency>
 
     <dependency>
       <groupId>io.dropwizard.modules</groupId>

openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java

@@ -121,6 +121,7 @@ import org.openmetadata.service.socket.OpenMetadataAssetServlet;
 import org.openmetadata.service.socket.SocketAddressFilter;
 import org.openmetadata.service.socket.WebSocketManager;
 import org.openmetadata.service.util.MicrometerBundleSingleton;
+import org.openmetadata.service.util.jdbi.DatabaseAuthenticationProviderFactory;
 import org.openmetadata.service.workflows.searchIndex.SearchIndexEvent;
 import org.quartz.SchedulerException;
 
@@ -273,6 +274,16 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
   }
 
   private Jdbi createAndSetupJDBI(Environment environment, DataSourceFactory dbFactory) {
+    // Check for db auth providers.
+    DatabaseAuthenticationProviderFactory.get(dbFactory.getUrl())
+        .ifPresent(
+            databaseAuthenticationProvider -> {
+              String token =
+                  databaseAuthenticationProvider.authenticate(
+                      dbFactory.getUrl(), dbFactory.getUser(), dbFactory.getPassword());
+              dbFactory.setPassword(token);
+            });
+
     Jdbi jdbi = new JdbiFactory().build(environment, dbFactory, "database");
     SqlLogger sqlLogger =
         new SqlLogger() {

openmetadata-service/src/main/java/org/openmetadata/service/util/TablesInitializer.java

@@ -53,6 +53,7 @@ import org.openmetadata.service.search.IndexUtil;
 import org.openmetadata.service.search.SearchClient;
 import org.openmetadata.service.search.SearchIndexDefinition;
 import org.openmetadata.service.secrets.SecretsManagerFactory;
+import org.openmetadata.service.util.jdbi.DatabaseAuthenticationProviderFactory;
 
 public final class TablesInitializer {
   private static final String DEBUG_MODE_ENABLED = "debug_mode";
@@ -183,9 +184,21 @@ public final class TablesInitializer {
     if (dataSourceFactory == null) {
       throw new RuntimeException("No database in config file");
     }
+
+    // Check for db auth providers.
+    DatabaseAuthenticationProviderFactory.get(dataSourceFactory.getUrl())
+        .ifPresent(
+            databaseAuthenticationProvider -> {
+              String token =
+                  databaseAuthenticationProvider.authenticate(
+                      dataSourceFactory.getUrl(), dataSourceFactory.getUser(), dataSourceFactory.getPassword());
+              dataSourceFactory.setPassword(token);
+            });
+
     String jdbcUrl = dataSourceFactory.getUrl();
     String user = dataSourceFactory.getUser();
     String password = dataSourceFactory.getPassword();
+
     boolean disableValidateOnMigrate = commandLine.hasOption(DISABLE_VALIDATE_ON_MIGRATE);
     if (disableValidateOnMigrate) {
       printToConsoleInDebug("Disabling validation on schema migrate");

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/AwsRdsDatabaseAuthenticationProvider.java

@@ -0,0 +1,86 @@
+package org.openmetadata.service.util.jdbi;
+
+import java.io.UnsupportedEncodingException;
+import java.net.MalformedURLException;
+import java.net.URI;
+import java.net.URL;
+import java.net.URLDecoder;
+import java.nio.charset.StandardCharsets;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+import org.jetbrains.annotations.NotNull;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.rds.RdsUtilities;
+import software.amazon.awssdk.services.rds.model.GenerateAuthenticationTokenRequest;
+
+/**
+ * {@link DatabaseAuthenticationProvider} implementation for AWS RDS IAM Auth.
+ *
+ * @see <a href="https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.Enabling.html"></a>
+ */
+public class AwsRdsDatabaseAuthenticationProvider implements DatabaseAuthenticationProvider {
+
+  public static final String AWS_REGION = "awsRegion";
+  public static final String ALLOW_PUBLIC_KEY_RETRIEVAL = "allowPublicKeyRetrieval";
+  public static final String PROTOCOL = "https://";
+
+  @Override
+  public String authenticate(String jdbcUrl, String username, String password) {
+    // !!
+    try {
+      // Prepare
+      URI uri = URI.create(PROTOCOL + removeProtocolFrom(jdbcUrl));
+      Map<String, String> queryParams = parseQueryParams(uri.toURL());
+
+      // Set
+      String awsRegion = queryParams.get(AWS_REGION);
+      String allowPublicKeyRetrieval = queryParams.get(ALLOW_PUBLIC_KEY_RETRIEVAL);
+
+      // Validate
+      Objects.requireNonNull(awsRegion, "Parameter `awsRegion` shall be provided in the jdbc url.");
+      Objects.requireNonNull(
+          allowPublicKeyRetrieval, "Parameter `allowPublicKeyRetrieval` shall be provided in the jdbc url.");
+
+      // Prepare request
+      GenerateAuthenticationTokenRequest request =
+          GenerateAuthenticationTokenRequest.builder()
+              .credentialsProvider(DefaultCredentialsProvider.create())
+              .hostname(uri.getHost())
+              .port(uri.getPort())
+              .username(username)
+              .build();
+
+      // Return token
+      return RdsUtilities.builder().region(Region.of(awsRegion)).build().generateAuthenticationToken(request);
+
+    } catch (MalformedURLException | UnsupportedEncodingException e) {
+      // Throw
+      throw new DatabaseAuthenticationProviderException(e);
+    }
+  }
+
+  @NotNull
+  private static String removeProtocolFrom(String jdbcUrl) {
+    return jdbcUrl.substring(jdbcUrl.indexOf("://") + 3);
+  }
+
+  private Map<String, String> parseQueryParams(URL url) throws UnsupportedEncodingException {
+    // Prepare
+    Map<String, String> query_pairs = new LinkedHashMap<>();
+    String query = url.getQuery();
+    String[] pairs = query.split("&");
+
+    // Loop
+    for (String pair : pairs) {
+      int idx = pair.indexOf("=");
+      // Add
+      query_pairs.put(
+          URLDecoder.decode(pair.substring(0, idx), StandardCharsets.UTF_8),
+          URLDecoder.decode(pair.substring(idx + 1), StandardCharsets.UTF_8));
+    }
+    // Return
+    return query_pairs;
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/DatabaseAuthenticationProvider.java

@@ -0,0 +1,18 @@
+package org.openmetadata.service.util.jdbi;
+
+/**
+ * Database authentication provider is the main interface responsible for all implementation that requires additional
+ * authentication steps required by the database in order to authorize a user to be able to operate on it.
+ *
+ * <p>For example if a jdbc url requires to retrieve and authorized token this interface shall be implemented to
+ * retrieve the token.
+ */
+public interface DatabaseAuthenticationProvider {
+
+  /**
+   * Authenticate a user for the given jdbc url.
+   *
+   * @return authorization token
+   */
+  String authenticate(String jdbcUrl, String username, String password);
+}

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/DatabaseAuthenticationProviderException.java

@@ -0,0 +1,23 @@
+package org.openmetadata.service.util.jdbi;
+
+/** Database authentication provider exception responsible to all generic exception thrown by this layer. */
+public class DatabaseAuthenticationProviderException extends RuntimeException {
+  public DatabaseAuthenticationProviderException() {}
+
+  public DatabaseAuthenticationProviderException(String message) {
+    super(message);
+  }
+
+  public DatabaseAuthenticationProviderException(String message, Throwable cause) {
+    super(message, cause);
+  }
+
+  public DatabaseAuthenticationProviderException(Throwable cause) {
+    super(cause);
+  }
+
+  public DatabaseAuthenticationProviderException(
+      String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
+    super(message, cause, enableSuppression, writableStackTrace);
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/DatabaseAuthenticationProviderFactory.java

@@ -0,0 +1,27 @@
+package org.openmetadata.service.util.jdbi;
+
+import java.util.Optional;
+
+/** Factory class for {@link DatabaseAuthenticationProvider}. */
+public class DatabaseAuthenticationProviderFactory {
+  /** C'tor */
+  private DatabaseAuthenticationProviderFactory() {}
+
+  /**
+   * Get auth provider based on the given jdbc url.
+   *
+   * @param jdbcURL the jdbc url.
+   * @return instance of {@link DatabaseAuthenticationProvider}.
+   */
+  public static Optional<DatabaseAuthenticationProvider> get(String jdbcURL) {
+    // Check
+    if (jdbcURL.contains(AwsRdsDatabaseAuthenticationProvider.AWS_REGION)
+        && jdbcURL.contains(AwsRdsDatabaseAuthenticationProvider.ALLOW_PUBLIC_KEY_RETRIEVAL)) {
+      // Return AWS RDS Auth provider
+      return Optional.of(new AwsRdsDatabaseAuthenticationProvider());
+    }
+
+    // Return empty
+    return Optional.empty();
+  }
+}