Initial.

Author: claudevdm

.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": 5
+  "modification": 6
 }

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

@@ -72,6 +72,8 @@ public Schema configurationSchema() {
         // readQuery
         .addNullableField("partitionColumn", FieldType.STRING)
         .addNullableField("partitions", FieldType.INT16)
+        .addNullableField("longLowerBound", FieldType.INT64) // For numeric bounds
+        .addNullableField("longUpperBound", FieldType.INT64) // For numeric bounds
         .addNullableField("maxConnections", FieldType.INT16)
         .addNullableField("driverJars", FieldType.STRING)
         .addNullableField("writeBatchSize", FieldType.INT64)
@@ -139,6 +141,18 @@ public PCollection<Row> expand(PBegin input) {
               readRows = readRows.withNumPartitions(partitions);
             }
 
+            if (config.getSchema().hasField("longLowerBound")
+                && config.getSchema().hasField("longUpperBound")
+                && config.getInt64("longLowerBound") != null
+                && config.getInt64("longUpperBound") != null) {
+              readRows =
+                  ((JdbcIO.ReadWithPartitions<Row, Long>) readRows)
+                      .withLowerBound(
+                          Preconditions.checkStateNotNull(config.getInt64("longLowerBound")))
+                      .withUpperBound(
+                          Preconditions.checkStateNotNull(config.getInt64("longUpperBound")));
+            }
+
             @Nullable Short fetchSize = config.getInt16("fetchSize");
             if (fetchSize != null) {
               readRows = readRows.withFetchSize(fetchSize);

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

@@ -131,6 +131,31 @@ public void testPartitionedReadWithExplicitSchema() {
     pipeline.run();
   }
 
+  @Test
+  public void testPartitionedReadWithExplicitLongBounds() {
+    JdbcSchemaIOProvider provider = new JdbcSchemaIOProvider();
+
+    Row config =
+        Row.withSchema(provider.configurationSchema())
+            .withFieldValue("driverClassName", DATA_SOURCE_CONFIGURATION.getDriverClassName().get())
+            .withFieldValue("jdbcUrl", DATA_SOURCE_CONFIGURATION.getUrl().get())
+            .withFieldValue("username", "")
+            .withFieldValue("password", "")
+            .withFieldValue("partitionColumn", "id")
+            .withFieldValue("partitions", (short) 5)
+            .withFieldValue("longLowerBound", 0L)
+            .withFieldValue("longUpperBound", (long) EXPECTED_ROW_COUNT / 2)
+            .build();
+
+    JdbcSchemaIOProvider.JdbcSchemaIO schemaIO =
+        provider.from(READ_TABLE_NAME, config, Schema.builder().build());
+
+    PCollection<Row> output = pipeline.apply(schemaIO.buildReader());
+    Long expected = Long.valueOf(EXPECTED_ROW_COUNT / 2 + 1);
+    PAssert.that(output.apply(Count.globally())).containsInAnyOrder(expected);
+    pipeline.run();
+  }
+
   @Test
   public void testReadWithExplicitSchema() {
     JdbcSchemaIOProvider provider = new JdbcSchemaIOProvider();

sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py

@@ -288,6 +288,26 @@ def test_xlang_jdbc_write_read(self, database):
 
       assert_that(result, equal_to(expected_rows))
 
+    # Try the same read using the partitioned reader code path.
+    # with explicit partition bounds.
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      result = (
+          p
+          | 'Partitioned read from jdbc' >> ReadFromJdbc(
+              table_name=table_name,
+              partition_column='f_id',
+              partitions=3,
+              lower_bound=0,
+              upper_bound=2,
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath']))
+
+      assert_that(result, equal_to(expected_rows[:3]))
+
   @parameterized.expand(['postgres', 'mysql'])
   def test_xlang_jdbc_read_with_explicit_schema(self, database):
     if self.containers[database] is None:

sdks/python/apache_beam/io/jdbc.py

@@ -131,6 +131,8 @@ def default_io_expansion_service(classpath=None):
      ('autosharding', typing.Optional[bool]),
      ('partition_column', typing.Optional[str]),
      ('partitions', typing.Optional[np.int16]),
+     ('long_lower_bound', typing.Optional[np.int64]),
+     ('long_upper_bound', typing.Optional[np.int64]),
      ('max_connections', typing.Optional[np.int16]),
      ('driver_jars', typing.Optional[str]),
      ('write_batch_size', typing.Optional[np.int64])],
@@ -250,7 +252,9 @@ def __init__(
                             max_connections=max_connections,
                             driver_jars=driver_jars,
                             partitions=None,
-                            partition_column=None)),
+                            partition_column=None,
+                            long_lower_bound=None,
+                            long_upper_bound=None)),
                 dataSchema=None),
         ),
         expansion_service or default_io_expansion_service(classpath),
@@ -301,6 +305,8 @@ def __init__(
       fetch_size=None,
       partition_column=None,
       partitions=None,
+      lower_bound=None,
+      upper_bound=None,
       connection_properties=None,
       connection_init_sqls=None,
       max_connections=None,
@@ -381,7 +387,9 @@ def __init__(
                             max_connections=max_connections,
                             driver_jars=driver_jars,
                             partition_column=partition_column,
-                            partitions=partitions)),
+                            partitions=partitions,
+                            long_lower_bound=lower_bound,
+                            long_upper_bound=upper_bound)),
                 dataSchema=dataSchema),
         ),
         expansion_service or default_io_expansion_service(classpath),