Fix handling of empty arrays, json arrays and jsonb arrays. (#2169)

* Fix handling of empty arrays, json arrays and jsonb arrays. * Add support for INTERVAL types. --------- Co-authored-by: Claude <[email protected]>
GoogleCloudPlatform · Feb 7, 2025 · dae14d7 · dae14d7
1 parent 33e76c6
commit dae14d7
Show file tree

Hide file tree

Showing 2 changed files with 207 additions and 4 deletions.
diff --git a/...ream-to-sql/src/main/java/com/google/cloud/teleport/v2/utils/DatastreamToPostgresDML.java b/...ream-to-sql/src/main/java/com/google/cloud/teleport/v2/utils/DatastreamToPostgresDML.java
@@ -106,15 +106,18 @@ public String cleanDataTypeValueSql(
           return getNullValueSql();
         }
         break;
+      case "INTERVAL":
+        return convertJsonToPostgresInterval(columnValue, columnName);
     }
+
     // Arrays in Postgres are prefixed with underscore e.g. _INT4 for integer array.
     if (dataType.startsWith("_")) {
-      return convertJsonToPostgresArray(columnValue);
+      return convertJsonToPostgresArray(columnValue, dataType.toUpperCase(), columnName);
     }
     return columnValue;
   }
 
-  private String convertJsonToPostgresArray(String jsonValue) {
+  public String convertJsonToPostgresInterval(String jsonValue, String columnName) {
     if (jsonValue == null || jsonValue.equals("''") || jsonValue.equals("")) {
       return getNullValueSql();
     }
@@ -123,8 +126,39 @@ private String convertJsonToPostgresArray(String jsonValue) {
       ObjectMapper mapper = new ObjectMapper();
       JsonNode rootNode = mapper.readTree(jsonValue);
 
+      if (!rootNode.isObject()
+          || !rootNode.has("months")
+          || !rootNode.has("hours")
+          || !rootNode.has("micros")) {
+        LOG.warn("Invalid interval format for column {}, value: {}", columnName, jsonValue);
+        return getNullValueSql();
+      }
+
+      int months = rootNode.get("months").asInt();
+      int hours = rootNode.get("hours").asInt();
+      double seconds = rootNode.get("micros").asLong() / 1_000_000.0;
+
+      // Build the ISO 8601 string
+      String intervalStr = String.format("P%dMT%dH%.6fS", months, hours, seconds);
+
+      return "'" + intervalStr + "'";
+
+    } catch (JsonProcessingException e) {
+      LOG.error("Error parsing JSON interval: {}", jsonValue, e);
+      return getNullValueSql();
+    }
+  }
+
+  private String convertJsonToPostgresArray(String jsonValue, String dataType, String columnName) {
+    if (jsonValue == null || jsonValue.equals("''") || jsonValue.equals("")) {
+      return getNullValueSql();
+    }
+
+    try {
+      ObjectMapper mapper = new ObjectMapper();
+      JsonNode rootNode = mapper.readTree(jsonValue);
       if (!(rootNode.isObject() && rootNode.has("nestedArray"))) {
-        LOG.warn("Empty array: {}", jsonValue);
+        LOG.warn("Null array for column {}, value {}", columnName, jsonValue);
         return getNullValueSql();
       }
 
@@ -146,7 +180,20 @@ private String convertJsonToPostgresArray(String jsonValue) {
           }
         }
       }
-      return "ARRAY[" + String.join(",", elements) + "]";
+      if (elements.isEmpty()) {
+        // Use array literal for empty arrays otherwise type inferencing fails.
+        return "'{}'";
+      }
+      String arrayStatement = "ARRAY[" + String.join(",", elements) + "]";
+      if (dataType.equals("_JSON")) {
+        // Cast string array to json array.
+        return arrayStatement + "::json[]";
+      }
+      if (dataType.equals("_JSONB")) {
+        // Cast string array to jsonb array.
+        return arrayStatement + "::jsonb[]";
+      }
+      return arrayStatement;
     } catch (JsonProcessingException e) {
       LOG.error("Error parsing JSON array: {}", jsonValue);
       return getNullValueSql();

diff --git a/...tastream-to-sql/src/test/java/com/google/cloud/teleport/v2/utils/DatastreamToDMLTest.java b/...tastream-to-sql/src/test/java/com/google/cloud/teleport/v2/utils/DatastreamToDMLTest.java
@@ -154,6 +154,162 @@ public void testTextArrayTypeCoercion() {
     assertEquals(expectedInt, actualInt);
   }
 
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts an
+   * empty array into the correct PostgreSQL empty array literal '{}'.
+   */
+  @Test
+  public void testEmptyArray() {
+    String arrayJson = "{\"empty_array\": {\"nestedArray\": []}}";
+    JsonNode rowObj = getRowObj(arrayJson);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("empty_array", "_TEXT"); // Use a generic array type; could be any array
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'{}'";
+
+    String actual = dml.getValueSql(rowObj, "empty_array", tableSchema);
+
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a
+   * JSONB array into the correct PostgreSQL array syntax with type casting.
+   */
+  @Test
+  public void testJsonbArray() {
+    String arrayJson =
+        "{\"jsonb_array\": {"
+            + "\"nestedArray\": ["
+            + "  {\"nestedArray\": null, \"elementValue\": {\"a\": 1, \"b\": \"test\"}},"
+            + "  {\"nestedArray\": null, \"elementValue\": {\"c\": true}}"
+            + "], \"elementValue\": null}}";
+    JsonNode rowObj = getRowObj(arrayJson);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("jsonb_array", "_JSONB"); // Explicitly specify JSONB array type
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "ARRAY[{\"a\":1,\"b\":\"test\"},{\"c\":true}]::jsonb[]";
+
+    String actual = dml.getValueSql(rowObj, "jsonb_array", tableSchema);
+
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a JSON
+   * array into the correct PostgreSQL array syntax with type casting.
+   */
+  @Test
+  public void testJsonArray() {
+    String arrayJson =
+        "{\"json_array\": {"
+            + "\"nestedArray\": ["
+            + "  {\"nestedArray\": null, \"elementValue\": {\"x\": 10, \"y\": \"abc\"}},"
+            + "  {\"nestedArray\": null, \"elementValue\": {\"z\": false}}"
+            + "], \"elementValue\": null}}";
+    JsonNode rowObj = getRowObj(arrayJson);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("json_array", "_JSON"); // Explicitly specify JSON array type
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "ARRAY[{\"x\":10,\"y\":\"abc\"},{\"z\":false}]::json[]";
+
+    String actual = dml.getValueSql(rowObj, "json_array", tableSchema);
+
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a JSON
+   * INTERVAL array into the correct PostgreSQL syntax.
+   */
+  @Test
+  public void testValidInterval() {
+    String json = "{\"interval_field\": {\"months\": 1, \"hours\": 2, \"micros\": 3000000}}";
+    JsonNode rowObj = getRowObj(json);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("interval_field", "INTERVAL");
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'P1MT2H3.000000S'";
+    String actual = dml.getValueSql(rowObj, "interval_field", tableSchema);
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a JSON
+   * INTERVAL array into the correct PostgreSQL syntax.
+   */
+  @Test
+  public void testOnlyMonths() {
+    String json = "{\"interval_field\": {\"months\": 12, \"hours\": 0, \"micros\": 0}}";
+    JsonNode rowObj = getRowObj(json);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("interval_field", "INTERVAL");
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'P12MT0H0.000000S'";
+    String actual = dml.getValueSql(rowObj, "interval_field", tableSchema);
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a JSON
+   * INTERVAL array into the correct PostgreSQL syntax.
+   */
+  @Test
+  public void testOnlyHours() {
+    String json = "{\"interval_field\": {\"months\": 0, \"hours\": 5, \"micros\": 0}}";
+    JsonNode rowObj = getRowObj(json);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("interval_field", "INTERVAL");
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'P0MT5H0.000000S'";
+    String actual = dml.getValueSql(rowObj, "interval_field", tableSchema);
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a JSON
+   * INTERVAL array into the correct PostgreSQL syntax.
+   */
+  @Test
+  public void testOnlyMicros() {
+    String json = "{\"interval_field\": {\"months\": 0, \"hours\": 0, \"micros\": 123456}}";
+    JsonNode rowObj = getRowObj(json);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("interval_field", "INTERVAL");
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'P0MT0H0.123456S'";
+    String actual = dml.getValueSql(rowObj, "interval_field", tableSchema);
+    assertEquals(expected, actual);
+  }
+
+  /**
+   * Test whether {@link DatastreamToPostgresDML#getValueSql(JsonNode, String, Map)} converts a JSON
+   * INTERVAL array into the correct PostgreSQL syntax.
+   */
+  @Test
+  public void testLargeMicros() {
+    String json = "{\"interval_field\": {\"months\": 0, \"hours\": 0, \"micros\": 999999999}}";
+    JsonNode rowObj = getRowObj(json);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("interval_field", "INTERVAL");
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'P0MT0H999.999999S'";
+    String actual = dml.getValueSql(rowObj, "interval_field", tableSchema);
+    assertEquals(expected, actual);
+  }
+
+  @Test
+  public void testZeroValues() {
+    String json = "{\"interval_field\": {\"months\": 0, \"hours\": 0, \"micros\": 0}}";
+    JsonNode rowObj = getRowObj(json);
+    Map<String, String> tableSchema = new HashMap<>();
+    tableSchema.put("interval_field", "INTERVAL");
+    DatastreamToPostgresDML dml = DatastreamToPostgresDML.of(null);
+    String expected = "'P0MT0H0.000000S'";
+    String actual = dml.getValueSql(rowObj, "interval_field", tableSchema);
+    assertEquals(expected, actual);
+  }
+
   /**
    * Test whether {@link DatastreamToDML#getTargetSchemaName} converts the Oracle schema into the
    * correct Postgres schema.