Add ability to set schema in JdbcIO.java and jdbc.py Read.

Author: claudevdm

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());
+      // Use the provided schema if it's not null, otherwise infer it
+      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();

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,11 @@ public PCollection<Row> expand(PBegin input) {
               readRows = readRows.withDisableAutoCommit(disableAutoCommit);
             }
 
+            // If a schema was provided, use it
+            if (dataSchema != null) {
+              readRows = readRows.withSchema(dataSchema);
+            }
+
             return input.apply(readRows);
           }
         }

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(
@@ -305,7 +306,7 @@ def __init__(
       driver_jars=None,
       expansion_service=None,
       classpath=None,
-  ):
+      schema=None):
     """
     Initializes a read operation from Jdbc.
 
@@ -343,6 +344,14 @@ def __init__(
                       driver.
     """
     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 +376,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),
     )