added defaultQueryTimeout for jdbc io (#35376)

Author: liferoad

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

@@ -491,6 +491,9 @@ public abstract static class DataSourceConfiguration implements Serializable {
     @Pure
     abstract @Nullable ValueProvider<Integer> getMaxConnections();
 
+    @Pure
+    abstract @Nullable ValueProvider<Integer> getQueryTimeout();
+
     @Pure
     abstract @Nullable ClassLoader getDriverClassLoader();
 
@@ -520,6 +523,8 @@ abstract Builder setConnectionInitSqls(
 
       abstract Builder setMaxConnections(ValueProvider<@Nullable Integer> maxConnections);
 
+      abstract Builder setQueryTimeout(ValueProvider<@Nullable Integer> queryTimeout);
+
       abstract Builder setDriverClassLoader(ClassLoader driverClassLoader);
 
       abstract Builder setDriverJars(ValueProvider<String> driverJars);
@@ -622,6 +627,17 @@ public DataSourceConfiguration withMaxConnections(
       return builder().setMaxConnections(maxConnections).build();
     }
 
+    /** Sets the default query timeout that will be used for connections created by this source. */
+    public DataSourceConfiguration withQueryTimeout(Integer queryTimeout) {
+      checkArgument(queryTimeout != null, "queryTimeout can not be null");
+      return withQueryTimeout(ValueProvider.StaticValueProvider.of(queryTimeout));
+    }
+
+    /** Same as {@link #withQueryTimeout(Integer)} but accepting a ValueProvider. */
+    public DataSourceConfiguration withQueryTimeout(ValueProvider<@Nullable Integer> queryTimeout) {
+      return builder().setQueryTimeout(queryTimeout).build();
+    }
+
     /**
      * Sets the class loader instance to be used to load the JDBC driver. If not specified, the
      * default class loader is used.
@@ -657,6 +673,7 @@ void populateDisplayData(DisplayData.Builder builder) {
         builder.addIfNotNull(DisplayData.item("jdbcUrl", getUrl()));
         builder.addIfNotNull(DisplayData.item("username", getUsername()));
         builder.addIfNotNull(DisplayData.item("driverJars", getDriverJars()));
+        builder.addIfNotNull(DisplayData.item("queryTimeout", getQueryTimeout()));
       }
     }
 
@@ -701,6 +718,12 @@ public DataSource buildDatasource() {
             basicDataSource.setMaxTotal(maxConnections);
           }
         }
+        if (getQueryTimeout() != null) {
+          Integer queryTimeout = getQueryTimeout().get();
+          if (queryTimeout != null) {
+            basicDataSource.setDefaultQueryTimeout(queryTimeout);
+          }
+        }
         if (getDriverClassLoader() != null) {
           basicDataSource.setDriverClassLoader(getDriverClassLoader());
         }