Repository: apache/beam
Pull request #34128: Add explicit schema support to JdbcIO read and xlang transform.
Author: login: claudevdm
Labels: name: python
name: java
name: io
name: jdbc
name: Next Action: Reviewers
Created: 2025-02-28T21:51:35Z | Merged: 2025-03-03T20:13:58Z

- Added withSchema() to JdbcIO.ReadRows and JdbcIO.ReadWithPartitions in Java
- Wired schema through JdbcSchemaIOProvider
- Added schema support to xlang jcbc.py ReadFromJdbc
- Refactored xlang tests to initializing containers once at the class level instead of per test
- Added tests for custom SQL statements in Python (there was a todo in the code)

Postcommit xlang test run https://github.com/apache/beam/actions/runs/13635975481/job/38114752052?pr=34128
------------------------

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

 - [ ] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
 - [ ] Update `CHANGES.md` with noteworthy changes.
 - [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).

See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://github.com/apache/beam/blob/master/CONTRIBUTING.md#make-the-reviewers-job-easier).

To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)

GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[![Build python source distribution and wheels](https://github.com/apache/beam/actions/workflows/build_wheels.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[![Python tests](https://github.com/apache/beam/actions/workflows/python_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[![Java tests](https://github.com/apache/beam/actions/workflows/java_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
[![Go tests](https://github.com/apache/beam/actions/workflows/go_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)

See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.



Diff:
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
index a31745754d0d..a6f0e1b417ea 100644
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
@@ -745,6 +745,9 @@ public abstract static class ReadRows extends PTransform<PBegin, PCollection<Row
     @Pure
     abstract boolean getDisableAutoCommit();
 
+    @Pure
+    abstract @Nullable Schema getSchema();
+
     abstract Builder toBuilder();
 
     @AutoValue.Builder
@@ -762,6 +765,8 @@ abstract Builder setDataSourceProviderFn(
 
       abstract Builder setDisableAutoCommit(boolean disableAutoCommit);
 
+      abstract Builder setSchema(@Nullable Schema schema);
+
       abstract ReadRows build();
     }
 
@@ -789,6 +794,10 @@ public ReadRows withStatementPreparator(StatementPreparator statementPreparator)
       return toBuilder().setStatementPreparator(statementPreparator).build();
     }
 
+    public ReadRows withSchema(Schema schema) {
+      return toBuilder().setSchema(schema).build();
+    }
+
     /**
      * This method is used to set the size of the data that is going to be fetched and loaded in
      * memory per every database call. Please refer to: {@link java.sql.Statement#setFetchSize(int)}
@@ -830,7 +839,14 @@ public PCollection<Row> expand(PBegin input) {
               getDataSourceProviderFn(),
               "withDataSourceConfiguration() or withDataSourceProviderFn() is required");
 
-      Schema schema = inferBeamSchema(dataSourceProviderFn.apply(null), query.get());
+      // Don't infer schema if explicitly provided.
+      Schema schema;
+      if (getSchema() != null) {
+        schema = getSchema();
+      } else {
+        schema = inferBeamSchema(dataSourceProviderFn.apply(null), query.get());
+      }
+
       PCollection<Row> rows =
           input.apply(
               JdbcIO.<Row>read()
@@ -1292,6 +1308,9 @@ public abstract static class ReadWithPartitions<T, PartitionColumnT>
     @Pure
     abstract boolean getUseBeamSchema();
 
+    @Pure
+    abstract @Nullable Schema getSchema();
+
     @Pure
     abstract @Nullable PartitionColumnT getLowerBound();
 
@@ -1333,6 +1352,8 @@ abstract Builder<T, PartitionColumnT> setDataSourceProviderFn(
 
       abstract Builder<T, PartitionColumnT> setUseBeamSchema(boolean useBeamSchema);
 
+      abstract Builder setSchema(@Nullable Schema schema);
+
       abstract Builder<T, PartitionColumnT> setFetchSize(int fetchSize);
 
       abstract Builder<T, PartitionColumnT> setTable(String tableName);
@@ -1424,6 +1445,10 @@ public ReadWithPartitions<T, PartitionColumnT> withTable(String tableName) {
       return toBuilder().setTable(tableName).build();
     }
 
+    public ReadWithPartitions<T, PartitionColumnT> withSchema(Schema schema) {
+      return toBuilder().setSchema(schema).build();
+    }
+
     private static final int EQUAL = 0;
 
     @Override
@@ -1532,8 +1557,11 @@ public KV<Long, KV<PartitionColumnT, PartitionColumnT>> apply(
       Schema schema = null;
       if (getUseBeamSchema()) {
         schema =
-            ReadRows.inferBeamSchema(
-                dataSourceProviderFn.apply(null), String.format("SELECT * FROM %s", getTable()));
+            getSchema() != null
+                ? getSchema()
+                : ReadRows.inferBeamSchema(
+                    dataSourceProviderFn.apply(null),
+                    String.format("SELECT * FROM %s", getTable()));
         rowMapper = (RowMapper<T>) SchemaUtil.BeamRowMapper.of(schema);
       } else {
         rowMapper = getRowMapper();
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java
index 23221042938b..b9c8f2fad15d 100644
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java
@@ -84,7 +84,7 @@ public Schema configurationSchema() {
    */
   @Override
   public JdbcSchemaIO from(String location, Row configuration, @Nullable Schema dataSchema) {
-    return new JdbcSchemaIO(location, configuration);
+    return new JdbcSchemaIO(location, configuration, dataSchema);
   }
 
   @Override
@@ -101,10 +101,12 @@ public PCollection.IsBounded isBounded() {
   static class JdbcSchemaIO implements SchemaIO, Serializable {
     protected final Row config;
     protected final String location;
+    protected final @Nullable Schema dataSchema;
 
-    JdbcSchemaIO(String location, Row config) {
+    JdbcSchemaIO(String location, Row config, @Nullable Schema dataSchema) {
       this.config = config;
       this.location = location;
+      this.dataSchema = dataSchema;
     }
 
     @Override
@@ -147,6 +149,10 @@ public PCollection<Row> expand(PBegin input) {
               readRows = readRows.withDisableAutoCommit(disableAutoCommit);
             }
 
+            if (dataSchema != null) {
+              readRows = readRows.withSchema(dataSchema);
+            }
+
             return input.apply(readRows);
           } else {
 
@@ -175,6 +181,9 @@ public PCollection<Row> expand(PBegin input) {
               readRows = readRows.withDisableAutoCommit(disableAutoCommit);
             }
 
+            if (dataSchema != null) {
+              readRows = readRows.withSchema(dataSchema);
+            }
             return input.apply(readRows);
           }
         }
diff --git a/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcIOTest.java b/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcIOTest.java
index bc75aace9f77..34e7a31ecc0c 100644
--- a/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcIOTest.java
+++ b/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcIOTest.java
@@ -470,6 +470,51 @@ public void testReadWithSchema() {
     pipeline.run();
   }
 
+  @Test
+  public void testReadRowsWithExplicitSchema() {
+    Schema customSchema =
+        Schema.of(
+            Schema.Field.of("CUSTOMER_NAME", Schema.FieldType.STRING).withNullable(true),
+            Schema.Field.of("CUSTOMER_ID", Schema.FieldType.INT64).withNullable(true));
+
+    PCollection<Row> rows =
+        pipeline.apply(
+            JdbcIO.readRows()
+                .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION)
+                .withQuery(String.format("select name,id from %s where name = ?", READ_TABLE_NAME))
+                .withStatementPreparator(
+                    preparedStatement -> preparedStatement.setString(1, TestRow.getNameForSeed(1)))
+                .withSchema(customSchema));
+
+    assertEquals(customSchema, rows.getSchema());
+
+    PCollection<Row> output = rows.apply(Select.fieldNames("CUSTOMER_NAME", "CUSTOMER_ID"));
+    PAssert.that(output)
+        .containsInAnyOrder(
+            ImmutableList.of(Row.withSchema(customSchema).addValues("Testval1", 1L).build()));
+
+    pipeline.run();
+  }
+
+  @Test
+  @SuppressWarnings({"UnusedVariable"})
+  public void testIncompatibleSchemaThrowsError() {
+    Schema incompatibleSchema =
+        Schema.of(
+            Schema.Field.of("WRONG_TYPE_NAME", Schema.FieldType.INT64),
+            Schema.Field.of("WRONG_TYPE_ID", Schema.FieldType.STRING));
+
+    Pipeline pipeline = Pipeline.create();
+    pipeline.apply(
+        JdbcIO.readRows()
+            .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION)
+            .withQuery(String.format("select name,id from %s limit 10", READ_TABLE_NAME))
+            .withSchema(incompatibleSchema));
+
+    PipelineExecutionException exception =
+        assertThrows(PipelineExecutionException.class, () -> pipeline.run().waitUntilFinish());
+  }
+
   @Test
   public void testReadWithPartitions() {
     PCollection<TestRow> rows =
@@ -486,6 +531,32 @@ public void testReadWithPartitions() {
     pipeline.run();
   }
 
+  @Test
+  public void testReadWithPartitionsWithExplicitSchema() {
+    Schema customSchema =
+        Schema.of(
+            Schema.Field.of("CUSTOMER_NAME", Schema.FieldType.STRING).withNullable(true),
+            Schema.Field.of("CUSTOMER_ID", Schema.FieldType.INT32).withNullable(true));
+
+    PCollection<Row> rows =
+        pipeline.apply(
+            JdbcIO.<Row>readWithPartitions()
+                .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION)
+                .withTable(String.format("(select name,id from %s) as subq", READ_TABLE_NAME))
+                .withNumPartitions(5)
+                .withPartitionColumn("id")
+                .withLowerBound(0L)
+                .withUpperBound(1000L)
+                .withRowOutput()
+                .withSchema(customSchema));
+
+    assertEquals(customSchema, rows.getSchema());
+
+    PAssert.thatSingleton(rows.apply("Count All", Count.globally())).isEqualTo(1000L);
+
+    pipeline.run();
+  }
+
   @Test
   public void testReadWithPartitionsBySubqery() {
     PCollection<TestRow> rows =
diff --git a/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProviderTest.java b/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProviderTest.java
index 193a1f0c3477..3a8b83b992e2 100644
--- a/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProviderTest.java
+++ b/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProviderTest.java
@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.io.jdbc;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 
 import java.sql.Connection;
 import java.sql.PreparedStatement;
@@ -85,6 +86,92 @@ public void testPartitionedRead() {
     pipeline.run();
   }
 
+  @Test
+  public void testPartitionedReadWithExplicitSchema() {
+    JdbcSchemaIOProvider provider = new JdbcSchemaIOProvider();
+
+    Schema customSchema =
+        Schema.of(
+            Schema.Field.of("CUSTOMER_NAME", Schema.FieldType.STRING).withNullable(true),
+            Schema.Field.of("CUSTOMER_ID", Schema.FieldType.INT32).withNullable(true));
+
+    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) 10)
+            .build();
+
+    JdbcSchemaIOProvider.JdbcSchemaIO schemaIO =
+        provider.from(
+            String.format("(select name,id from %s) as subq", READ_TABLE_NAME),
+            config,
+            customSchema);
+
+    PCollection<Row> output = pipeline.apply(schemaIO.buildReader());
+
+    assertEquals(customSchema, output.getSchema());
+
+    Long expected = Long.valueOf(EXPECTED_ROW_COUNT);
+    PAssert.that(output.apply(Count.globally())).containsInAnyOrder(expected);
+
+    PAssert.that(output)
+        .satisfies(
+            rows -> {
+              for (Row row : rows) {
+                assertNotNull(row.getString("CUSTOMER_NAME"));
+                assertNotNull(row.getInt32("CUSTOMER_ID"));
+              }
+              return null;
+            });
+
+    pipeline.run();
+  }
+
+  @Test
+  public void testReadWithExplicitSchema() {
+    JdbcSchemaIOProvider provider = new JdbcSchemaIOProvider();
+
+    Schema customSchema =
+        Schema.of(
+            Schema.Field.of("CUSTOMER_NAME", Schema.FieldType.STRING).withNullable(true),
+            Schema.Field.of("CUSTOMER_ID", Schema.FieldType.INT32).withNullable(true));
+
+    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("readQuery", "SELECT name, id FROM " + READ_TABLE_NAME)
+            .build();
+
+    JdbcSchemaIOProvider.JdbcSchemaIO schemaIO =
+        provider.from(READ_TABLE_NAME, config, customSchema);
+
+    PCollection<Row> output = pipeline.apply(schemaIO.buildReader());
+
+    assertEquals(customSchema, output.getSchema());
+
+    Long expected = Long.valueOf(EXPECTED_ROW_COUNT);
+    PAssert.that(output.apply(Count.globally())).containsInAnyOrder(expected);
+
+    PAssert.that(output)
+        .satisfies(
+            rows -> {
+              for (Row row : rows) {
+                assertNotNull(row.getString("CUSTOMER_NAME"));
+                assertNotNull(row.getInt32("CUSTOMER_ID"));
+              }
+              return null;
+            });
+
+    pipeline.run();
+  }
+
   // This test shouldn't work because we only support numeric and datetime columns and we are trying
   // to use a string column as our partition source.
   @Test
diff --git a/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py b/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
index 38a405c2d331..d6abe4fdd2a5 100644
--- a/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
@@ -23,8 +23,6 @@
 import typing
 import unittest
 from decimal import Decimal
-from typing import Callable
-from typing import Union
 
 from parameterized import parameterized
 
@@ -56,6 +54,8 @@
 # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
 
 ROW_COUNT = 10
+POSTGRES_BINARY_TYPE = ('BYTEA', 'BYTEA')
+MYSQL_BINARY_TYPE = ('BINARY(10)', 'VARBINARY(10)')
 
 JdbcTestRow = typing.NamedTuple(
     "JdbcTestRow",
@@ -66,6 +66,27 @@
 )
 coders.registry.register_coder(JdbcTestRow, coders.RowCoder)
 
+CustomSchemaRow = typing.NamedTuple(
+    "CustomSchemaRow",
+    [
+        ("renamed_id", int),
+        ("renamed_float", float),
+        ("renamed_char", str),
+        ("renamed_varchar", str),
+        ("renamed_bytes", bytes),
+        ("renamed_varbytes", bytes),
+        ("renamed_timestamp", Timestamp),
+        ("renamed_decimal", Decimal),
+        ("renamed_date", datetime.date),
+        ("renamed_time", datetime.time),
+    ],
+)
+coders.registry.register_coder(CustomSchemaRow, coders.RowCoder)
+
+SimpleRow = typing.NamedTuple(
+    "SimpleRow", [("id", int), ("name", str), ("value", float)])
+coders.registry.register_coder(SimpleRow, coders.RowCoder)
+
 
 @unittest.skipIf(sqlalchemy is None, 'sql alchemy package is not installed.')
 @unittest.skipIf(
@@ -91,55 +112,73 @@ class CrossLanguageJdbcIOTest(unittest.TestCase):
           'com.mysql.cj.jdbc.Driver')
   }
 
-  def _setUpTestCase(
-      self,
-      container_init: Callable[[], Union[PostgresContainer, MySqlContainer]],
-      db_string: str,
-      driver: str):
-    # This method is not the normal setUp from unittest, because the test has
-    # beem parameterized. The setup then needs extra parameters to initialize.
-    self.start_db_container(retries=3, container_init=container_init)
-    self.engine = sqlalchemy.create_engine(self.db.get_connection_url())
-    self.username = 'test'
-    self.password = 'test'
-    self.host = self.db.get_container_host_ip()
-    self.port = self.db.get_exposed_port(self.db.port_to_expose)
-    self.database_name = 'test'
-    self.driver_class_name = driver
-    self.jdbc_url = 'jdbc:{}://{}:{}/{}'.format(
-        db_string, self.host, self.port, self.database_name)
-
-  def tearDown(self):
-    # Sometimes stopping the container raises ReadTimeout. We can ignore it
-    # here to avoid the test failure.
-    try:
-      self.db.stop()
-    except:  # pylint: disable=bare-except
-      logging.error('Could not stop the postgreSQL container.')
+  @classmethod
+  def setUpClass(cls):
+    cls.containers = {}
+    cls.engines = {}
+    cls.jdbc_configs = {}
 
-  @parameterized.expand(['postgres', 'mysql'])
-  def test_xlang_jdbc_write_read(self, database):
-    container_init, classpath, db_string, driver = (
-        CrossLanguageJdbcIOTest.DB_CONTAINER_CLASSPATH_STRING[database])
-    self._setUpTestCase(container_init, db_string, driver)
-    table_name = 'jdbc_external_test'
-    if database == 'postgres':
-      # postgres does not have BINARY and VARBINARY type, use equvalent.
-      binary_type = ('BYTEA', 'BYTEA')
-    else:
-      binary_type = ('BINARY(10)', 'VARBINARY(10)')
-
-    with self.engine.begin() as connection:
-      connection.execute(
-          sqlalchemy.text(
-              "CREATE TABLE IF NOT EXISTS {}".format(table_name) +
-              "(f_id INTEGER, " + "f_float DOUBLE PRECISION, " +
-              "f_char CHAR(10), " + "f_varchar VARCHAR(10), " +
-              f"f_bytes {binary_type[0]}, " + f"f_varbytes {binary_type[1]}, " +
-              "f_timestamp TIMESTAMP(3), " + "f_decimal DECIMAL(10, 2), " +
-              "f_date DATE, " + "f_time TIME(3))"))
-
-    inserted_rows = [
+    for db_type, db_data in cls.DB_CONTAINER_CLASSPATH_STRING.items():
+      container = cls.start_container(db_data.container_fn)
+      cls.containers[db_type] = container
+
+      cls.engines[db_type] = sqlalchemy.create_engine(
+          container.get_connection_url())
+
+      cls.jdbc_configs[db_type] = {
+          'username': 'test',
+          'password': 'test',
+          'host': container.get_container_host_ip(),
+          'port': container.get_exposed_port(container.port_to_expose),
+          'database_name': 'test',
+          'driver_class_name': db_data.connector,
+          'classpath': db_data.classpath,
+          'jdbc_url': (
+              f'jdbc:{db_data.db_string}://{container.get_container_host_ip()}:'
+              f'{container.get_exposed_port(container.port_to_expose)}/test'),
+          'binary_type': POSTGRES_BINARY_TYPE
+          if db_type == 'postgres' else MYSQL_BINARY_TYPE
+      }
+
+  @classmethod
+  def tearDownClass(cls):
+    for db_type, container in cls.containers.items():
+      if container:
+        # Sometimes stopping the container raises ReadTimeout. We can ignore it
+        # here to avoid the test failure.
+        try:
+          container.stop()
+        except Exception:  # pylint: disable=broad-except
+          logging.warning("Could not stop %s container", db_type)
+
+  @classmethod
+  def start_container(cls, container_init, max_retries=3):
+    # Creating a container with testcontainers sometimes raises ReadTimeout
+    # error. In java there are 2 retries set by default.
+    for attempt in range(max_retries):
+      try:
+        container = container_init()
+        container.start()
+        return container
+      except Exception:  # pylint: disable=broad-except
+        if attempt == max_retries - 1:
+          logging.error(
+              'Failed to initialize container after %s attempts', max_retries)
+          raise
+
+  def create_test_table(self, connection, table_name, database):
+    binary_type = self.jdbc_configs[database]['binary_type']
+    connection.execute(
+        sqlalchemy.text(
+            f"CREATE TABLE IF NOT EXISTS {table_name}" +
+            "(f_id INTEGER, f_float DOUBLE PRECISION, " +
+            "f_char CHAR(10), f_varchar VARCHAR(10), " +
+            f"f_bytes {binary_type[0]}, f_varbytes {binary_type[1]}, " +
+            "f_timestamp TIMESTAMP(3), f_decimal DECIMAL(10, 2), " +
+            "f_date DATE, f_time TIME(3))"))
+
+  def generate_test_data(self, count):
+    return [
         JdbcTestRow(
             i,
             i + 0.1,
@@ -147,24 +186,31 @@ def test_xlang_jdbc_write_read(self, database):
             f'Test{i}',
             f'Test{i}'.encode(),
             f'Test{i}'.encode(),
-            # In alignment with Java Instant which supports milli precision.
             Timestamp.of(seconds=round(time.time(), 3)),
-            # Test both positive and negative numbers.
             Decimal(f'{i-1}.23'),
-            # Test both date before or after EPOCH
             datetime.date(1969 + i, i % 12 + 1, i % 31 + 1),
             datetime.time(i % 24, i % 60, i % 60, (i * 1000) % 1_000_000))
-        for i in range(ROW_COUNT)
+        for i in range(count)
     ]
-    expected_row = []
-    for row in inserted_rows:
+
+  @parameterized.expand(['postgres', 'mysql'])
+  def test_xlang_jdbc_write_read(self, database):
+    table_name = f"jdbc_write_read_test_{database}"
+
+    with self.engines[database].begin() as connection:
+      self.create_test_table(connection, table_name, database)
+
+    test_rows = self.generate_test_data(ROW_COUNT)
+
+    expected_rows = []
+    for row in test_rows:
       f_char = row.f_char + ' ' * (10 - len(row.f_char))
+      f_bytes = row.f_bytes
+
       if database != 'postgres':
-        # padding expected results
         f_bytes = row.f_bytes + b'\0' * (10 - len(row.f_bytes))
-      else:
-        f_bytes = row.f_bytes
-      expected_row.append(
+
+      expected_rows.append(
           JdbcTestRow(
               row.f_id,
               row.f_float,
@@ -177,20 +223,20 @@ def test_xlang_jdbc_write_read(self, database):
               row.f_date,
               row.f_time))
 
+    config = self.jdbc_configs[database]
+
     with TestPipeline() as p:
       p.not_use_test_runner_api = True
       _ = (
           p
-          | beam.Create(inserted_rows).with_output_types(JdbcTestRow)
-          # TODO(https://github.com/apache/beam/issues/20446) Add test with
-          # overridden write_statement
+          | beam.Create(test_rows).with_output_types(JdbcTestRow)
           | 'Write to jdbc' >> WriteToJdbc(
               table_name=table_name,
-              driver_class_name=self.driver_class_name,
-              jdbc_url=self.jdbc_url,
-              username=self.username,
-              password=self.password,
-              classpath=classpath,
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath'],
           ))
 
     # Register MillisInstant logical type to override the mapping from Timestamp
@@ -201,17 +247,15 @@ def test_xlang_jdbc_write_read(self, database):
       p.not_use_test_runner_api = True
       result = (
           p
-          # TODO(https://github.com/apache/beam/issues/20446) Add test with
-          # overridden read_query
           | 'Read from jdbc' >> ReadFromJdbc(
               table_name=table_name,
-              driver_class_name=self.driver_class_name,
-              jdbc_url=self.jdbc_url,
-              username=self.username,
-              password=self.password,
-              classpath=classpath))
+              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_row))
+      assert_that(result, equal_to(expected_rows))
 
     # Try the same read using the partitioned reader code path.
     # Outputs should be the same.
@@ -223,26 +267,177 @@ def test_xlang_jdbc_write_read(self, database):
               table_name=table_name,
               partition_column='f_id',
               partitions=3,
-              driver_class_name=self.driver_class_name,
-              jdbc_url=self.jdbc_url,
-              username=self.username,
-              password=self.password,
-              classpath=classpath))
-
-      assert_that(result, equal_to(expected_row))
-
-  # Creating a container with testcontainers sometimes raises ReadTimeout
-  # error. In java there are 2 retries set by default.
-  def start_db_container(self, retries, container_init):
-    for i in range(retries):
-      try:
-        self.db = container_init()
-        self.db.start()
-        break
-      except Exception as e:  # pylint: disable=bare-except
-        if i == retries - 1:
-          logging.error('Unable to initialize database container.')
-          raise e
+              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))
+
+  @parameterized.expand(['postgres', 'mysql'])
+  def test_xlang_jdbc_read_with_explicit_schema(self, database):
+    if self.containers[database] is None:
+      self.skipTest(f"{database} container could not be initialized")
+
+    table_name = f"jdbc_schema_test_{database}"
+
+    with self.engines[database].begin() as connection:
+      self.create_test_table(connection, table_name, database)
+
+    test_rows = self.generate_test_data(ROW_COUNT)
+
+    expected_rows = []
+    for row in test_rows:
+      f_char = row.f_char
+      f_bytes = row.f_bytes
+
+      if database != 'postgres':
+        f_bytes = row.f_bytes + b'\0' * (10 - len(row.f_bytes))
+
+      expected_rows.append(
+          CustomSchemaRow(
+              row.f_id,
+              row.f_float,
+              f_char,
+              row.f_varchar,
+              f_bytes,
+              row.f_bytes,
+              row.f_timestamp,
+              row.f_decimal,
+              row.f_date,
+              row.f_time))
+
+    def custom_row_equals(expected, actual):
+      return (
+          expected.renamed_id == actual.renamed_id and
+          expected.renamed_float == actual.renamed_float and
+          expected.renamed_char.rstrip() == actual.renamed_char.rstrip() and
+          expected.renamed_varchar == actual.renamed_varchar and
+          expected.renamed_bytes == actual.renamed_bytes and
+          expected.renamed_timestamp == actual.renamed_timestamp and
+          expected.renamed_decimal == actual.renamed_decimal and
+          expected.renamed_date == actual.renamed_date and
+          expected.renamed_time == actual.renamed_time)
+
+    config = self.jdbc_configs[database]
+
+    # Run write pipeline
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      _ = (
+          p
+          | beam.Create(test_rows).with_output_types(JdbcTestRow)
+          | 'Write to jdbc' >> WriteToJdbc(
+              table_name=table_name,
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath'],
+          ))
+
+    # Register MillisInstant logical type to override the mapping from Timestamp
+    # originally handled by MicrosInstant.
+    LogicalType.register_logical_type(MillisInstant)
+
+    # Run read pipeline with custom schema
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      result = (
+          p
+          | 'Read from jdbc with schema' >> ReadFromJdbc(
+              table_name=table_name,
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath'],
+              schema=CustomSchemaRow))
+
+      assert_that(result, equal_to(expected_rows, equals_fn=custom_row_equals))
+
+  @parameterized.expand(['postgres', 'mysql'])
+  def test_xlang_jdbc_custom_statements(self, database):
+    # Skip if container wasn't initialized
+    if self.containers[database] is None:
+      self.skipTest(f"{database} container could not be initialized")
+
+    # Create a simple table for this test
+    table_name = f"jdbc_custom_statements_{database}"
+
+    with self.engines[database].begin() as connection:
+      connection.execute(
+          sqlalchemy.text(
+              f"CREATE TABLE IF NOT EXISTS {table_name}" +
+              "(id INTEGER, name VARCHAR(50), value DOUBLE PRECISION)"))
+
+    test_rows = [
+        SimpleRow(1, "Item1", 10.5),
+        SimpleRow(2, "Item2", 20.75),
+        SimpleRow(3, "Item3", 30.25),
+        SimpleRow(4, "Item4", 40.0),
+        SimpleRow(5, "Item5", 50.5)
+    ]
+
+    config = self.jdbc_configs[database]
+
+    write_statement = f"INSERT INTO {table_name} (id, name, value) VALUES \
+        (?, ?, ?)"
+
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      _ = (
+          p
+          | beam.Create(test_rows).with_output_types(SimpleRow)
+          | 'Write with custom statement' >> WriteToJdbc(
+              table_name="",
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath'],
+              statement=write_statement))
+
+    # Schema inference fails when there is a WHERE clause, so we pass explicit
+    # schema.
+    read_query = f"SELECT id, name, value FROM {table_name} WHERE value > 25.0"
+    expected_filtered_rows = [row for row in test_rows if row.value > 25.0]
+
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      result = (
+          p
+          | 'Read with custom query' >> ReadFromJdbc(
+              table_name="",
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath'],
+              query=read_query,
+              schema=SimpleRow))
+
+      assert_that(result, equal_to(expected_filtered_rows))
+
+    # JdbcIO#readWithPartitions requires custom queries to be passed as a
+    # wrapped subquery to table_name.
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      result = (
+          p
+          | 'Read with custom query' >> ReadFromJdbc(
+              table_name=f"({read_query}) as subq",
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=config['password'],
+              classpath=config['classpath'],
+              partition_column="id",
+              schema=SimpleRow))
+
+      # revert
+      assert_that(result, equal_to(test_rows))
 
 
 if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/io/jdbc.py b/sdks/python/apache_beam/io/jdbc.py
index 11570680a2f3..32ce16b358f8 100644
--- a/sdks/python/apache_beam/io/jdbc.py
+++ b/sdks/python/apache_beam/io/jdbc.py
@@ -114,7 +114,8 @@ def default_io_expansion_service(classpath=None):
 
 JdbcConfigSchema = typing.NamedTuple(
     'JdbcConfigSchema',
-    [('location', str), ('config', bytes)],
+    [('location', str), ('config', bytes),
+     ('dataSchema', typing.Optional[bytes])],
 )
 
 Config = typing.NamedTuple(
@@ -249,7 +250,8 @@ def __init__(
                             max_connections=max_connections,
                             driver_jars=driver_jars,
                             partitions=None,
-                            partition_column=None))),
+                            partition_column=None)),
+                dataSchema=None),
         ),
         expansion_service or default_io_expansion_service(classpath),
     )
@@ -305,7 +307,7 @@ def __init__(
       driver_jars=None,
       expansion_service=None,
       classpath=None,
-  ):
+      schema=None):
     """
     Initializes a read operation from Jdbc.
 
@@ -341,8 +343,20 @@ def __init__(
                       package (e.g. "org.postgresql:postgresql:42.3.1").
                       By default, this argument includes a Postgres SQL JDBC
                       driver.
+    :param schema: Optional custom schema for the returned rows. If provided,
+                   this should be a NamedTuple type that defines the structure
+                   of the output PCollection elements. This bypasses automatic
+                   schema inference during pipeline construction.
     """
     classpath = classpath or DEFAULT_JDBC_CLASSPATH
+
+    dataSchema = None
+    if schema is not None:
+      # Convert Python schema to Beam Schema proto
+      schema_proto = typing_to_runner_api(schema).row_type.schema
+      # Serialize the proto to bytes for transmission
+      dataSchema = schema_proto.SerializeToString()
+
     super().__init__(
         self.URN,
         NamedTupleBasedPayloadBuilder(
@@ -367,7 +381,8 @@ def __init__(
                             max_connections=max_connections,
                             driver_jars=driver_jars,
                             partition_column=partition_column,
-                            partitions=partitions))),
+                            partitions=partitions)),
+                dataSchema=dataSchema),
         ),
         expansion_service or default_io_expansion_service(classpath),
     )


Reviews:
user: login: Abacn
body: Thanks!

Just to confirm, have you been able to verify it on Dataflow runner? e.g. submit a jdbc read pipeline locally where it does not have access to database; and Dataflow workers have access and the pipeline runs successfully
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-03T19:39:28Z
user: login: claudevdm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T19:55:43Z
user: login: Abacn
body: SGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-03T20:13:42Z

Review comments:
diff_hunk: @@ -84,7 +84,7 @@ public Schema configurationSchema() {
    */
   @Override
   public JdbcSchemaIO from(String location, Row configuration, @Nullable Schema dataSchema) {
-    return new JdbcSchemaIO(location, configuration);
+    return new JdbcSchemaIO(location, configuration, dataSchema);
path: sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java
user: login: Abacn
body: I see, previously the dataSchema parameter wasn't passed down stream
created_at: 2025-03-03T19:38:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -84,7 +84,7 @@ public Schema configurationSchema() {
    */
   @Override
   public JdbcSchemaIO from(String location, Row configuration, @Nullable Schema dataSchema) {
-    return new JdbcSchemaIO(location, configuration);
+    return new JdbcSchemaIO(location, configuration, dataSchema);
path: sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcSchemaIOProvider.java
user: login: claudevdm
body: Yeah and in the JdbcIO#readRows and readWithPartitions there schema case wasnt handled.

I also need to add the ability to explicitly pass the partition lower and upper bounds through xlang, otherwise that part will also happen at pipeline construction time. Will do in a followup later this week.
created_at: 2025-03-03T19:55:43Z
author_association: COLLABORATOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-02T02:41:16Z
body: ## [Codecov](https://app.codecov.io/gh/apache/beam/pull/34128?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
Attention: Patch coverage is `0%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 59.30%. Comparing base [(`f5ed586`)](https://app.codecov.io/gh/apache/beam/commit/f5ed5862bd97d4bc20c3b98f146eca3ec406e500?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) to head [(`6832b61`)](https://app.codecov.io/gh/apache/beam/commit/6832b6136e73e9f327b86385902acd6f99abcb89?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
> Report is 58 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/apache/beam/pull/34128?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [sdks/python/apache\_beam/io/jdbc.py](https://app.codecov.io/gh/apache/beam/pull/34128?src=pr&el=tree&filepath=sdks%2Fpython%2Fapache_beam%2Fio%2Fjdbc.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vaW8vamRiYy5weQ==) | 0.00% | [4 Missing :warning: ](https://app.codecov.io/gh/apache/beam/pull/34128?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master   #34128      +/-   ##
============================================
+ Coverage     59.25%   59.30%   +0.04%     
  Complexity     3272     3272              
============================================
  Files          1164     1166       +2     
  Lines        178325   178633     +308     
  Branches       3413     3393      -20     
============================================
+ Hits         105675   105936     +261     
- Misses        69250    69302      +52     
+ Partials       3400     3395       -5     
```

| [Flag](https://app.codecov.io/gh/apache/beam/pull/34128/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [python](https://app.codecov.io/gh/apache/beam/pull/34128/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `81.27% <0.00%> (+0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/beam/pull/34128?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).

author_association: NONE
user: login: claudevdm
created_at: 2025-03-03T17:58:51Z
body: addresses #23029
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-03-03T19:06:19Z
body: Assigning reviewers. If you would like to opt out of this review, comment `assign to next reviewer`:

R: @damccorm for label python.
R: @Abacn for label java.
R: @shunping for label io.

Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).
author_association: CONTRIBUTOR
user: login: claudevdm
created_at: 2025-03-03T19:59:18Z
body: > Thanks!
> 
> Just to confirm, have you been able to verify it on Dataflow runner? e.g. submit a jdbc read pipeline locally where it does not have access to database; and Dataflow workers have access and the pipeline runs successfully

I confirmed that infer schema doesn't happen during pipeline construction when the schema parameter is passed  :D . 

I didn't try to run on DataflowRunner because I am not sure how I am supposed to use the modified Java SDK expansion service. For example, I usually pass something like sdk_harness_container_image_overrides=.*java.*,gcr.io/cloud-dataflow/v1beta3/beam_java11_sdk:2.62.0, does that mean I need to build the Java SDK container etc?

If you look at the custom query/write statement tests, those schema inferences would file if we don't pass a custom schema. So it indirectly verifies the solution.
author_association: COLLABORATOR