From 860baf0e4fee8c0433839937160b078a517dcb13 Mon Sep 17 00:00:00 2001
From: Derrick Williams <derrickaw@google.com>
Date: Thu, 6 Aug 2026 18:09:34 +0000
Subject: [PATCH] add jdbc pipeline and other fixes

---
 sdks/python/apache_beam/typehints/schemas.py  |  1 +
 .../extended_tests/databases/postgres.yaml    | 41 ++++++++++++++++++-
 .../apache_beam/yaml/integration_tests.py     |  8 +++-
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py
index 2fd3c22e1e58..1861bc43b936 100644
--- a/sdks/python/apache_beam/typehints/schemas.py
+++ b/sdks/python/apache_beam/typehints/schemas.py
@@ -1420,6 +1420,7 @@ def argument(self):
 # TODO: A temporary fix for missing jdbc logical types.
 # See the discussion in https://github.com/apache/beam/issues/35738 for
 # more detail.
+@LogicalType._register_internal
 class JdbcDateType(LogicalType[datetime.date, MillisInstant, str]):
   """
   For internal use only; no backwards-compatibility guarantees.
diff --git a/sdks/python/apache_beam/yaml/extended_tests/databases/postgres.yaml b/sdks/python/apache_beam/yaml/extended_tests/databases/postgres.yaml
index 8957c782753e..ad06acb89f56 100644
--- a/sdks/python/apache_beam/yaml/extended_tests/databases/postgres.yaml
+++ b/sdks/python/apache_beam/yaml/extended_tests/databases/postgres.yaml
@@ -66,4 +66,43 @@ pipelines:
             elements:
               - {value: 123, rank: 0}
               - {value: 456, rank: 1}
-              - {value: 789, rank: 2}
\ No newline at end of file
+              - {value: 789, rank: 2}
+
+  # Postgres write pipeline using WriteToJdbc with date/time values
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {event_name: 'User Login', event_date: '2025-06-23', event_time: '10:00:00'}
+              - {event_name: 'Data Backup', event_date: '2025-06-23', event_time: '14:30:00'}
+              - {event_name: 'Maintenance Start', event_date: null, event_time: null}
+        - type: WriteToJdbc
+          config:
+            url: "{TEMP_DB}"
+            driver_class_name: "org.postgresql.Driver"
+            query: "INSERT INTO test_date_time (event_name, event_date, event_time) VALUES(?, CAST(? AS DATE), CAST(? AS TIME))"
+
+  # Postgres read pipeline using ReadFromJdbc with date/time values
+  - pipeline:
+      type: chain
+      transforms:
+        - type: ReadFromJdbc
+          config:
+            url: "{TEMP_DB}"
+            driver_class_name: "org.postgresql.Driver"
+            table: "test_date_time"
+        - type: MapToFields
+          config:
+            language: python
+            fields:
+              event_name: event_name
+              event_date: "str(event_date) if event_date is not None else None"
+              event_time: "str(event_time) if event_time is not None else None"
+        - type: AssertEqual
+          config:
+            elements:
+              - {event_name: 'User Login', event_date: '2025-06-23', event_time: '10:00:00'}
+              - {event_name: 'Data Backup', event_date: '2025-06-23', event_time: '14:30:00'}
+              - {event_name: 'Maintenance Start', event_date: null, event_time: null}
\ No newline at end of file
diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py
index c6d73df76e31..b0a07de19880 100644
--- a/sdks/python/apache_beam/yaml/integration_tests.py
+++ b/sdks/python/apache_beam/yaml/integration_tests.py
@@ -95,9 +95,9 @@ def get_impl(self):
 from apache_beam.io.gcp.spanner_wrapper import SpannerWrapper
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.utils import python_callable
+from apache_beam.yaml import conftest
 from apache_beam.yaml import yaml_provider
 from apache_beam.yaml import yaml_transform
-from apache_beam.yaml.conftest import yaml_test_files_dir
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -493,6 +493,10 @@ def temp_postgres_database():
         connection.execute(
             sqlalchemy.text(
                 "CREATE TABLE tmp_table (value INTEGER, rank INTEGER);"))
+        connection.execute(
+            sqlalchemy.text(
+                "CREATE TABLE test_date_time (event_name VARCHAR(100), event_date DATE, event_time TIME);"
+            ))
 
       # Construct the JDBC url for connections later on by tests
       jdbc_url = (
@@ -1029,7 +1033,7 @@ def parse_test_files(filepattern):
 
 # Dynamically create test methods from the tests directory.
 # yaml_test_files_dir comes from conftest.py and set by pytest_configure.
-_test_files_dir = yaml_test_files_dir
+_test_files_dir = conftest.yaml_test_files_dir
 _file_pattern = os.path.join(
     os.path.dirname(__file__), _test_files_dir, '*.yaml')
 parse_test_files(_file_pattern)