Add AWS Postgres Iam Auth jar to GMS (#6371)

* Add AWS Postgres Iam Auth jar to GMS * Added IAM support for Postgres * Setting IAM support for Postgres as false by default * format fix Co-authored-by: syed.javed <[email protected]>
datahub-project · Dec 5, 2022 · ffea930 · ffea930
1 parent fdcb731
commit ffea930
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 1 deletion.
diff --git a/build.gradle b/build.gradle
@@ -59,6 +59,8 @@ project.ext.externalDependency = [
     'awsGlueSchemaRegistrySerde': 'software.amazon.glue:schema-registry-serde:1.1.10',
     'awsMskIamAuth': 'software.amazon.msk:aws-msk-iam-auth:1.1.1',
     'awsSecretsManagerJdbc': 'com.amazonaws.secretsmanager:aws-secretsmanager-jdbc:1.0.8',
+    'awsPostgresIamAuth': 'software.amazon.jdbc:aws-advanced-jdbc-wrapper:1.0.0',
+    'awsRds':'software.amazon.awssdk:rds:2.18.24',
     'cacheApi' : 'javax.cache:cache-api:1.1.0',
     'commonsCli': 'commons-cli:commons-cli:1.5.0',
     'commonsIo': 'commons-io:commons-io:2.4',

diff --git a/docker/datahub-gms/env/docker.postgres.env b/docker/datahub-gms/env/docker.postgres.env
@@ -3,6 +3,9 @@ EBEAN_DATASOURCE_PASSWORD=datahub
 EBEAN_DATASOURCE_HOST=postgres:5432
 EBEAN_DATASOURCE_URL=jdbc:postgresql://postgres:5432/datahub
 EBEAN_DATASOURCE_DRIVER=org.postgresql.Driver
+# Uncomment EBEAN_POSTGRES_USE_AWS_IAM_AUTH below to add support for IAM authentication for Postgres.
+# Password is not required when accessing Postgres using IAM auth. It can be replaced by dummy password
+# EBEAN_POSTGRES_USE_AWS_IAM_AUTH=true
 KAFKA_BOOTSTRAP_SERVER=broker:29092
 KAFKA_SCHEMAREGISTRY_URL=http://schema-registry:8081
 ELASTICSEARCH_HOST=elasticsearch

diff --git a/metadata-service/factories/build.gradle b/metadata-service/factories/build.gradle
@@ -25,7 +25,8 @@ dependencies {
   compile externalDependency.springKafka
   compile externalDependency.springWeb
   compile project(':metadata-service:auth-ranger-impl')
-
+  implementation externalDependency.awsPostgresIamAuth
+  implementation externalDependency.awsRds
   annotationProcessor externalDependency.lombok
 
   compile spec.product.pegasus.restliSpringBridge

diff --git a/...actories/src/main/java/com/linkedin/gms/factory/common/LocalEbeanServerConfigFactory.java b/...actories/src/main/java/com/linkedin/gms/factory/common/LocalEbeanServerConfigFactory.java
@@ -6,6 +6,8 @@
 import io.ebean.datasource.DataSourceConfig;
 import io.ebean.datasource.DataSourcePoolListener;
 import java.sql.Connection;
+import java.util.HashMap;
+import java.util.Map;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
@@ -51,6 +53,9 @@ public class LocalEbeanServerConfigFactory {
   @Value("${ebean.autoCreateDdl:false}")
   private Boolean ebeanAutoCreate;
 
+  @Value("${ebean.postgresUseIamAuth:false}")
+  private Boolean postgresUseIamAuth;
+
   private DataSourcePoolListener getListenerToTrackCounts(String metricName) {
     final String counterName = "ebeans_connection_pool_size_" + metricName;
     return new DataSourcePoolListener() {
@@ -79,6 +84,12 @@ private DataSourceConfig buildDataSourceConfig(String dataSourceUrl, String data
     dataSourceConfig.setLeakTimeMinutes(ebeanLeakTimeMinutes);
     dataSourceConfig.setWaitTimeoutMillis(ebeanWaitTimeoutMillis);
     dataSourceConfig.setListener(getListenerToTrackCounts(dataSourceType));
+    // Adding IAM auth access for AWS Postgres
+    if (postgresUseIamAuth) {
+      Map<String, String> custom = new HashMap<>();
+      custom.put("wrapperPlugins", "iam");
+      dataSourceConfig.setCustomProperties(custom);
+    }
     return dataSourceConfig;
   }
 

diff --git a/metadata-service/factories/src/main/resources/application.yml b/metadata-service/factories/src/main/resources/application.yml
@@ -110,6 +110,7 @@ ebean:
   leakTimeMinutes: ${EBEAN_LEAK_TIME_MINUTES:15}
   waitTimeoutMillis: ${EBEAN_WAIT_TIMEOUT_MILLIS:1000}
   autoCreateDdl: ${EBEAN_AUTOCREATE:false}
+  postgresUseIamAuth: ${EBEAN_POSTGRES_USE_AWS_IAM_AUTH:false}
 
 # Only required if entityService.impl is cassandra
 cassandra: