Support Secret Manager in JdbcIO for Java, Python and YAML

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 1
+  "modification": 2
 }

.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "revision": 1
+  "revision": 2
 }

sdks/java/extensions/schemaio-expansion-service/build.gradle

@@ -66,6 +66,12 @@ dependencies {
     permitUnusedDeclared 'com.google.cloud.sql:postgres-socket-factory:1.25.0'
     implementation 'com.google.cloud.sql:mysql-socket-factory-connector-j-8:1.25.0'
     permitUnusedDeclared 'com.google.cloud.sql:mysql-socket-factory-connector-j-8:1.25.0'
+    implementation enforcedPlatform(library.java.google_cloud_platform_libraries_bom)
+    permitUnusedDeclared enforcedPlatform(library.java.google_cloud_platform_libraries_bom)
+    implementation library.java.google_cloud_secret_manager
+    permitUnusedDeclared library.java.google_cloud_secret_manager
+    implementation library.java.proto_google_cloud_secret_manager_v1
+    permitUnusedDeclared library.java.proto_google_cloud_secret_manager_v1
     testImplementation library.java.junit
     testImplementation library.java.mockito_core
     runtimeOnly ("org.xerial:sqlite-jdbc:3.49.1.0")

sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java

@@ -85,6 +85,7 @@
 import org.apache.beam.sdk.util.BackOff;
 import org.apache.beam.sdk.util.BackOffUtils;
 import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Secret;
 import org.apache.beam.sdk.util.Sleeper;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
@@ -320,6 +321,27 @@
  * since that risks duplicating records in the database, or failing due to primary key conflicts.
  * Consider using <a href="https://en.wikipedia.org/wiki/Merge_(SQL)">MERGE ("upsert")
  * statements</a> supported by your database instead.
+ *
+ * <h3>Using Secret Manager</h3>
+ *
+ * <p>Secret Manager is supported in both read and write operations to avoid storing sensitive
+ * credentials such as database passwords in plain text. You can configure Secret Manager on {@link
+ * DataSourceConfiguration} by specifying the secret manager provider using {@link
+ * DataSourceConfiguration#withSecretManager(String)} (e.g. {@code "GoogleCloudSecretManager"}) and
+ * providing the secret specification string in JSON format to {@link
+ * DataSourceConfiguration#withPassword(String)}.
+ *
+ * <p>For example for Google Cloud Secret Manager:
+ *
+ * <pre>{@code
+ * pipeline.apply(JdbcIO.<...>read()
+ *   .withDataSourceConfiguration(JdbcIO.DataSourceConfiguration.create(...)
+ *       ...
+ *       .withPassword("{\"name\": \"my-db-secret\", \"project\": \"my-project\"}")
+ *       .withSecretManager("GoogleCloudSecretManager"))
+ *   ...
+ * );
+ * }</pre>
  */
 @SuppressWarnings({
   "rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
@@ -503,6 +525,9 @@ public abstract static class DataSourceConfiguration implements Serializable {
     @Pure
     abstract @Nullable ValueProvider<String> getDriverJars();
 
+    @Pure
+    abstract @Nullable ValueProvider<@Nullable String> getSecretManager();
+
     @Pure
     abstract @Nullable DataSource getDataSource();
 
@@ -532,6 +557,8 @@ abstract Builder setConnectionInitSqls(
 
       abstract Builder setDriverJars(ValueProvider<String> driverJars);
 
+      abstract Builder setSecretManager(ValueProvider<@Nullable String> secretManager);
+
       abstract Builder setDataSource(@Nullable DataSource dataSource);
 
       abstract DataSourceConfiguration build();
@@ -571,10 +598,19 @@ public DataSourceConfiguration withUsername(ValueProvider<@Nullable String> user
       return builder().setUsername(username).build();
     }
 
+    /**
+     * Sets the database password.
+     *
+     * <p>You can specify a plain password string. Alternatively, if a secret manager is configured
+     * via {@link #withSecretManager(String)}, you can set this to a secret specification in JSON
+     * format (e.g. {@code "{\"name\": \"my-db-secret\", \"project\": \"my-project\"}"} for Google
+     * Cloud Secret Manager) that the secret manager uses to retrieve the password.
+     */
     public DataSourceConfiguration withPassword(@Nullable String password) {
       return withPassword(ValueProvider.StaticValueProvider.of(password));
     }
 
+    /** Same as {@link #withPassword(String)} but accepting a ValueProvider. */
     public DataSourceConfiguration withPassword(ValueProvider<@Nullable String> password) {
       return builder().setPassword(password).build();
     }
@@ -668,6 +704,28 @@ public DataSourceConfiguration withDriverJars(ValueProvider<String> driverJars)
       return builder().setDriverJars(driverJars).build();
     }
 
+    /**
+     * Sets the secret manager provider.
+     *
+     * <p>Currently supported options are:
+     *
+     * <ul>
+     *   <li>{@code "GoogleCloudSecretManager"}
+     *   <li>{@code "GoogleCloudHsmGeneratedSecretManager"}
+     * </ul>
+     *
+     * <p>If not set, no secret manager is used and the password is treated as a plain password.
+     */
+    public DataSourceConfiguration withSecretManager(@Nullable String secretManager) {
+      return withSecretManager(ValueProvider.StaticValueProvider.of(secretManager));
+    }
+
+    /** Same as {@link #withSecretManager(String)} but accepting a ValueProvider. */
+    public DataSourceConfiguration withSecretManager(
+        ValueProvider<@Nullable String> secretManager) {
+      return builder().setSecretManager(secretManager).build();
+    }
+
     void populateDisplayData(DisplayData.Builder builder) {
       if (getDataSource() != null) {
         builder.addIfNotNull(DisplayData.item("dataSource", getDataSource().getClass().getName()));
@@ -677,6 +735,7 @@ void populateDisplayData(DisplayData.Builder builder) {
         builder.addIfNotNull(DisplayData.item("username", getUsername()));
         builder.addIfNotNull(DisplayData.item("driverJars", getDriverJars()));
         builder.addIfNotNull(DisplayData.item("queryTimeout", getQueryTimeout()));
+        builder.addIfNotNull(DisplayData.item("secretManager", getSecretManager()));
       }
     }
 
@@ -689,6 +748,7 @@ public DataSource buildDatasource() {
         if (getUrl() != null) {
           basicDataSource.setUrl(getUrl().get());
         }
+        ValueProvider<@Nullable String> secretManagerProvider = getSecretManager();
         if (getUsername() != null) {
           @SuppressWarnings(
               "nullness") // this is actually nullable, but apache commons dbcp2 not annotated
@@ -701,6 +761,16 @@ public DataSource buildDatasource() {
               "nullness") // this is actually nullable, but apache commons dbcp2 not annotated
           @NonNull
           String password = getPassword().get();
+          if (password != null) {
+            String secretManager = null;
+            if (secretManagerProvider != null) {
+              secretManager = secretManagerProvider.get();
+            }
+            String fetched = Secret.fromJson(password, secretManager).getString(false);
+            if (fetched != null) {
+              password = fetched;
+            }
+          }
           basicDataSource.setPassword(password);
         }
         if (getConnectionProperties() != null) {

sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcReadSchemaTransformProvider.java

@@ -111,7 +111,24 @@ public String description() {
         + "      config:\n"
         + "        connectionProperties: \"characterEncoding=UTF-8;\"\n"
         + "        ...\n"
-        + "All properties should be semi-colon-delimited (e.g. \"key1=value1;key2=value2;\")\n";
+        + "All properties should be semi-colon-delimited (e.g. \"key1=value1;key2=value2;\")\n"
+        + "\n"
+        + "#### Using Secret Manager\n"
+        + "\n"
+        + "Secret Manager is supported to avoid storing sensitive credentials such as database passwords "
+        + "in plain text. You can configure `secret_manager` (e.g. `GoogleCloudSecretManager`) and provide the "
+        + "secret specification string in JSON format to `password`.\n"
+        + "\n"
+        + "For example, for Google Cloud Secret Manager: ::\n"
+        + "\n"
+        + "    - type: ReadFromJdbc\n"
+        + "      config:\n"
+        + "        jdbc_type: mysql\n"
+        + "        url: \"jdbc:mysql://my-host:3306/database\"\n"
+        + "        username: \"my-username\"\n"
+        + "        password: \"{\\\"name\\\": \\\"my-db-secret\\\", \\\"project\\\": \\\"my-project\\\"}\"\n"
+        + "        secret_manager: \"GoogleCloudSecretManager\"\n"
+        + "        query: \"SELECT * FROM table\"\n";
   }
 
   protected String inheritedDescription(
@@ -200,6 +217,11 @@ protected JdbcIO.DataSourceConfiguration dataSourceConfiguration() {
         dsConfig = dsConfig.withConnectionProperties(connectionProperties);
       }
 
+      String secretManager = config.getSecretManager();
+      if (secretManager != null) {
+        dsConfig = dsConfig.withSecretManager(secretManager);
+      }
+
       List<@org.checkerframework.checker.nullness.qual.Nullable String> initialSql =
           config.getConnectionInitSql();
       if (initialSql != null && initialSql.size() > 0) {
@@ -355,14 +377,20 @@ public abstract static class JdbcReadSchemaTransformConfiguration implements Ser
     @Nullable
     public abstract Boolean getOutputParallelization();
 
-    @SchemaFieldDescription("Password for the JDBC source.")
+    @SchemaFieldDescription(
+        "Password for the JDBC source. Can be specified as a plain password, or as a secret specification in JSON format if used with a secret manager.")
     @Nullable
     public abstract String getPassword();
 
     @SchemaFieldDescription("SQL query used to query the JDBC source.")
     @Nullable
     public abstract String getReadQuery();
 
+    @SchemaFieldDescription(
+        "Secret Manager to use for fetching secret values. Available options: 'GoogleCloudSecretManager', 'GoogleCloudHsmGeneratedSecretManager'. If not set, no secret manager is used and the password is treated as a plain password.")
+    @Nullable
+    public abstract String getSecretManager();
+
     @SchemaFieldDescription("Username for the JDBC source.")
     @Nullable
     public abstract String getUsername();
@@ -451,6 +479,8 @@ public abstract static class Builder {
 
       public abstract Builder setDriverJars(String value);
 
+      public abstract Builder setSecretManager(String value);
+
       public abstract JdbcReadSchemaTransformConfiguration build();
     }
   }

sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java

@@ -60,6 +60,7 @@ public Schema configurationSchema() {
         .addStringField("jdbcUrl")
         .addStringField("username")
         .addStringField("password")
+        .addNullableField("secretManager", FieldType.STRING)
         .addNullableField("connectionProperties", FieldType.STRING)
         .addNullableField("connectionInitSqls", FieldType.iterable(FieldType.STRING))
         .addNullableField("readQuery", FieldType.STRING)
@@ -221,6 +222,11 @@ protected JdbcIO.DataSourceConfiguration getDataSourceConfiguration() {
               .withUsername(config.getString("username"))
               .withPassword(config.getString("password"));
 
+      @Nullable String secretManager = config.getString("secretManager");
+      if (secretManager != null) {
+        dataSourceConfiguration = dataSourceConfiguration.withSecretManager(secretManager);
+      }
+
       @Nullable String connectionProperties = config.getString("connectionProperties");
       if (connectionProperties != null) {
         dataSourceConfiguration =

sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcWriteSchemaTransformProvider.java

@@ -115,7 +115,24 @@ public String description() {
         + "      config:\n"
         + "        connectionProperties: \"characterEncoding=UTF-8;\"\n"
         + "        ...\n"
-        + "All properties should be semi-colon-delimited (e.g. \"key1=value1;key2=value2;\")\n";
+        + "All properties should be semi-colon-delimited (e.g. \"key1=value1;key2=value2;\")\n"
+        + "\n"
+        + "#### Using Secret Manager\n"
+        + "\n"
+        + "Secret Manager is supported to avoid storing sensitive credentials such as database passwords "
+        + "in plain text. You can configure `secret_manager` (e.g. `GoogleCloudSecretManager`) and provide the "
+        + "secret specification string in JSON format to `password`.\n"
+        + "\n"
+        + "For example, for Google Cloud Secret Manager: ::\n"
+        + "\n"
+        + "    - type: WriteToJdbc\n"
+        + "      config:\n"
+        + "        jdbc_type: mysql\n"
+        + "        url: \"jdbc:mysql://my-host:3306/database\"\n"
+        + "        username: \"my-username\"\n"
+        + "        password: \"{\\\"name\\\": \\\"my-db-secret\\\", \\\"project\\\": \\\"my-project\\\"}\"\n"
+        + "        secret_manager: \"GoogleCloudSecretManager\"\n"
+        + "        query: \"INSERT INTO table VALUES(?, ?)\"\n";
   }
 
   protected String inheritedDescription(
@@ -204,6 +221,11 @@ protected JdbcIO.DataSourceConfiguration dataSourceConfiguration() {
         dsConfig = dsConfig.withConnectionProperties(connectionProperties);
       }
 
+      String secretManager = config.getSecretManager();
+      if (secretManager != null) {
+        dsConfig = dsConfig.withSecretManager(secretManager);
+      }
+
       List<@org.checkerframework.checker.nullness.qual.Nullable String> initialSql =
           config.getConnectionInitSql();
       if (initialSql != null && initialSql.size() > 0) {
@@ -340,10 +362,16 @@ public abstract static class JdbcWriteSchemaTransformConfiguration implements Se
     @Nullable
     public abstract String getLocation();
 
-    @SchemaFieldDescription("Password for the JDBC source.")
+    @SchemaFieldDescription(
+        "Password for the JDBC source. Can be specified as a plain password, or as a secret specification in JSON format if used with a secret manager.")
     @Nullable
     public abstract String getPassword();
 
+    @SchemaFieldDescription(
+        "Secret Manager to use for fetching secret values. Available options: 'GoogleCloudSecretManager', 'GoogleCloudHsmGeneratedSecretManager'. If not set, no secret manager is used and the password is treated as a plain password.")
+    @Nullable
+    public abstract String getSecretManager();
+
     @SchemaFieldDescription("Username for the JDBC source.")
     @Nullable
     public abstract String getUsername();
@@ -427,6 +455,8 @@ public abstract Builder setConnectionInitSql(
 
       public abstract Builder setBatchSize(Long value);
 
+      public abstract Builder setSecretManager(String value);
+
       public abstract JdbcWriteSchemaTransformConfiguration build();
     }
   }

sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcIOTest.java

@@ -32,6 +32,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -86,13 +87,15 @@
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.Wait;
+import org.apache.beam.sdk.util.Secret;
 import org.apache.beam.sdk.util.SerializableUtils;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.TypeDescriptors;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.commons.dbcp2.BasicDataSource;
 import org.apache.commons.dbcp2.PoolingDataSource;
 import org.apache.commons.lang3.StringUtils;
 import org.hamcrest.Description;
@@ -108,6 +111,7 @@
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.MockedStatic;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -195,6 +199,47 @@ public void testDataSourceConfigurationNullUsernameAndPassword() throws Exceptio
     }
   }
 
+  @Test
+  public void testDataSourceConfigurationPlainPasswordAndNullSecretManager() throws Exception {
+    String username = "sa";
+    String password = "my_plain_password";
+    JdbcIO.DataSourceConfiguration config =
+        DATA_SOURCE_CONFIGURATION
+            .withUsername(username)
+            .withPassword(password)
+            .withSecretManager((String) null);
+    DataSource dataSource = config.buildDatasource();
+    assertTrue(dataSource instanceof BasicDataSource);
+    assertEquals(password, ((BasicDataSource) dataSource).getPassword());
+    try (Connection conn = dataSource.getConnection()) {
+      assertTrue(conn.isValid(0));
+    }
+  }
+
+  @Test
+  public void testDataSourceConfigurationWithMockedSecretManager() {
+    String secretSpec = "{'name': 'my-db-secret'}";
+    String resolvedPassword = "my_fetched_secret_password";
+    String secretManager = "GoogleCloudSecretManager";
+
+    Secret mockSecret = mock(Secret.class);
+    when(mockSecret.getString(false)).thenReturn(resolvedPassword);
+
+    try (MockedStatic<Secret> mockedSecret = mockStatic(Secret.class)) {
+      mockedSecret.when(() -> Secret.fromJson(secretSpec, secretManager)).thenReturn(mockSecret);
+
+      JdbcIO.DataSourceConfiguration config =
+          DATA_SOURCE_CONFIGURATION
+              .withUsername("sa")
+              .withPassword(secretSpec)
+              .withSecretManager(secretManager);
+
+      DataSource dataSource = config.buildDatasource();
+      assertTrue(dataSource instanceof BasicDataSource);
+      assertEquals(resolvedPassword, ((BasicDataSource) dataSource).getPassword());
+    }
+  }
+
   @Test
   public void testSetConnectoinInitSqlFailWithDerbyDB() {
     String username = "sa";