With max num of records

Author: Amar3tto

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumReadSchemaTransformProvider.java

@@ -150,6 +150,18 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
               readTransform
                   .withMaxNumberOfRecords(testLimitRecords)
                   .withMaxTimeToRun(testLimitMilliseconds);
+        } else {
+          if (configuration.getMaxNumberOfRecords() != null) {
+            readTransform =
+              readTransform.withMaxNumberOfRecords(
+                configuration.getMaxNumberOfRecords());
+          }
+
+          if (configuration.getMaxTimeToRun() != null) {
+            readTransform =
+              readTransform.withMaxTimeToRun(
+                configuration.getMaxTimeToRun());
+          }
         }
 
         // TODO(pabloem): Database connection issues can be debugged here.
@@ -187,6 +199,14 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
   @DefaultSchema(AutoValueSchema.class)
   @AutoValue
   public abstract static class DebeziumReadSchemaTransformConfiguration {
+    @Nullable
+    @SchemaFieldDescription("Maximum number of records to read before stopping.")
+    public abstract Integer getMaxNumberOfRecords();
+
+    @Nullable
+    @SchemaFieldDescription("Maximum time in milliseconds to run before stopping.")
+    public abstract Long getMaxTimeToRun();
+
     @SchemaFieldDescription("Username used to connect to the source database.")
     public abstract String getUsername();
 

sdks/python/apache_beam/yaml/extended_tests/databases/debezium.yaml

@@ -31,6 +31,8 @@ pipelines:
             host: "{DEBEZIUM_DB[HOST]}"
             port: "{DEBEZIUM_DB[PORT]}"
             table: "{DEBEZIUM_DB[TABLE]}"
+            max_number_of_records: 2
+            max_time_to_run: 600000
             connection_properties:
               - "database.dbname={DEBEZIUM_DB[DATABASE]}"
               - "plugin.name=pgoutput"

sdks/python/apache_beam/yaml/standard_io.yaml

@@ -196,6 +196,8 @@
         port: 'port'
         table: 'table'
         connection_properties: 'debezium_connection_properties'
+        max_number_of_records: 'max_number_of_records'
+        max_time_to_run: 'max_time_to_run'
     underlying_provider:
       type: beamJar
       transforms: