Fix pylint; fix unsafe cast

apache · Feb 21, 2023 · db41b12 · db41b12
1 parent 5b6ef10
commit db41b12
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 13 deletions.
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
@@ -39,7 +39,6 @@
  * An implementation of {@link SchemaIOProvider} for reading and writing JSON payloads with {@link
  * JdbcIO}.
  */
-@SuppressWarnings({"unsafe"})
 @Internal
 @AutoService(SchemaIOProvider.class)
 public class JdbcSchemaIOProvider implements SchemaIOProvider {
@@ -123,8 +122,8 @@ public PCollection<Row> expand(PBegin input) {
                   ? config.getString("partitionColumn")
                   : null;
           if (partitionColumn != null) {
-            JdbcIO.ReadWithPartitions<?, ?> readRows =
-                JdbcIO.readWithPartitions()
+            JdbcIO.ReadWithPartitions<Row, ?> readRows =
+                JdbcIO.<Row>readWithPartitions()
                     .withDataSourceConfiguration(getDataSourceConfiguration())
                     .withTable(location)
                     .withPartitionColumn(partitionColumn)
@@ -133,10 +132,7 @@ public PCollection<Row> expand(PBegin input) {
             if (partitions != null) {
               readRows = readRows.withNumPartitions(partitions);
             }
-
-            // Need to do a cast here
-            return input.apply((JdbcIO.ReadWithPartitions<Row, ?>) readRows);
-
+            return input.apply(readRows);
           } else {
 
             @Nullable String readQuery = config.getString("readQuery");

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
@@ -201,10 +201,11 @@ def test_xlang_jdbc_write_read(self, database):
 
       assert_that(result, equal_to(expected_row))
 
-    # Try the same read using the partitioned reader code path. Outputs should be the same.
+    # Try the same read using the partitioned reader code path. 
+    # Outputs should be the same.
     with TestPipeline() as p:
-        p.not_use_test_runner_api = True
-        result = (
+      p.not_use_test_runner_api = True
+      result = (
           p
           | 'Partitioned read from jdbc' >> ReadFromJdbc(
           table_name=table_name,
@@ -216,7 +217,7 @@ def test_xlang_jdbc_write_read(self, database):
           password=self.password,
           classpath=classpath))
 
-        assert_that(result, equal_to(expected_row))
+      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.

diff --git a/sdks/python/apache_beam/io/jdbc.py b/sdks/python/apache_beam/io/jdbc.py
@@ -295,8 +295,10 @@ def __init__(
     :param query: sql query to be executed
     :param output_parallelization: is output parallelization on
     :param fetch_size: how many rows to fetch
-    :param partition_column: enable partitioned reads by splitting on this column.
-    :param partitions: override the default number of splits when using partition_column
+    :param partition_column: enable partitioned reads by splitting on this
+                             column
+    :param partitions: override the default number of splits when using
+                       partition_column
     :param connection_properties: properties of the jdbc connection
                                   passed as string with format
                                   [propertyName=property;]*