Bug fixes of handling null values and added missing IT

v2/spanner-custom-shard/src/main/java/com/custom/CustomTransformationWithCassandraForIT.java

@@ -46,6 +46,8 @@ public MigrationTransformationResponse toSourceRow(MigrationTransformationReques
       Map<String, Object> requestRow = request.getRequestRow();
       Map<String, Object> row = new HashMap<>();
       row.put("full_name", requestRow.get("first_name") + " " + requestRow.get("last_name"));
+      row.put("empty_string", "");
+      row.put("null_key", null);
       MigrationTransformationResponse response = new MigrationTransformationResponse(row, false);
       return response;
     }

v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/dml/CassandraDMLGenerator.java

@@ -174,16 +174,16 @@ private static DMLGeneratorResponse generateDMLResponse(
             .putAll(pkColumnNameValues)
             .putAll(columnNameValues)
             .build();
-    switch (modType) {
-      case "INSERT":
-      case "UPDATE":
-        return getUpsertStatementCQL(sourceTable.getName(), timestamp, allColumnNamesAndValues);
-      case "DELETE":
-        return getDeleteStatementCQL(sourceTable.getName(), timestamp, allColumnNamesAndValues);
-      default:
+    return switch (modType) {
+      case "INSERT", "UPDATE" -> getUpsertStatementCQL(
+          sourceTable.getName(), timestamp, allColumnNamesAndValues);
+      case "DELETE" -> getDeleteStatementCQL(
+          sourceTable.getName(), timestamp, allColumnNamesAndValues);
+      default -> {
         LOG.error("Unsupported modType: {} for table {}", modType, spannerTable.getName());
-        return new DMLGeneratorResponse("");
-    }
+        yield new DMLGeneratorResponse("");
+      }
+    };
   }
 
   /**
@@ -321,9 +321,11 @@ private static Map<String, PreparedStatementValueObject<?>> getColumnValues(
       if (customTransformColumns != null
           && customTransformColumns.contains(sourceColDef.getName())) {
         String cassandraType = sourceColDef.getType().getName().toLowerCase();
+        Object customValue = customTransformationResponse.get(colName);
         columnValue =
             PreparedStatementValueObject.create(
-                cassandraType, customTransformationResponse.get(colName));
+                cassandraType,
+                customValue == null ? CassandraTypeHandler.NullClass.INSTANCE : customValue);
         response.put(sourceColDef.getName(), columnValue);
         continue;
       }
@@ -398,9 +400,11 @@ private static Map<String, PreparedStatementValueObject<?>> getPkColumnValues(
           && customTransformColumns.contains(sourceColDef.getName())) {
         String cassandraType = sourceColDef.getType().getName().toLowerCase();
         String columnName = spannerColDef.getName();
+        Object customValue = customTransformationResponse.get(columnName);
         columnValue =
             PreparedStatementValueObject.create(
-                cassandraType, customTransformationResponse.get(columnName));
+                cassandraType,
+                customValue == null ? CassandraTypeHandler.NullClass.INSTANCE : customValue);
       } else if (keyValuesJson.has(spannerColumnName)) {
         columnValue =
             getMappedColumnValue(

v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/dbutils/dml/CassandraTypeHandler.java

@@ -41,6 +41,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.UUID;
 import org.apache.commons.lang3.BooleanUtils;
@@ -301,8 +302,7 @@ private static Object handleSpannerColumnType(
       String spannerType, String columnName, JSONObject valuesJson) {
     try {
       if (spannerType.contains("string")) {
-        String value = valuesJson.optString(columnName);
-        return value.isEmpty() ? null : value;
+        return valuesJson.optString(columnName, null);
       } else if (spannerType.contains("bytes")) {
         if (valuesJson.isNull(columnName)) {
           return null;
@@ -336,6 +336,9 @@ private static Object handleSpannerColumnType(
    */
   private static PreparedStatementValueObject<?> parseAndCastToCassandraType(
       String columnType, Object colValue) {
+    if (colValue == null) {
+      return null;
+    }
 
     if (columnType.startsWith("frozen<")) {
       return parseAndCastToCassandraType(extractInnerType(columnType), colValue);
@@ -531,7 +534,7 @@ private static LocalDate parseDate(Object colValue) {
    * @return a {@link List} containing parsed values, or an empty list if {@code colValue} is null
    */
   private static List<?> parseCassandraList(String columnType, JSONArray colValue) {
-    if (colValue == null) {
+    if (colValue.isEmpty()) {
       return Collections.emptyList();
     }
     String innerType = extractInnerType(columnType);
@@ -586,7 +589,7 @@ private static PreparedStatementValueObject<?> parseNestedType(String type, Obje
    * @return a {@link Set} containing parsed values, or an empty set if {@code colValue} is null
    */
   private static Set<?> parseCassandraSet(String columnType, JSONArray colValue) {
-    if (colValue == null) {
+    if (colValue.isEmpty()) {
       return Collections.emptySet();
     }
     String innerType = extractInnerType(columnType);
@@ -607,7 +610,7 @@ private static Set<?> parseCassandraSet(String columnType, JSONArray colValue) {
    *     null
    */
   private static Map<?, ?> parseCassandraMap(String columnType, JSONObject colValue) {
-    if (colValue == null) {
+    if (colValue.isEmpty()) {
       return Collections.emptyMap();
     }
     String[] keyValueTypes = extractKeyValueTypes(columnType);
@@ -651,11 +654,8 @@ public static PreparedStatementValueObject<?> getColumnValueByType(
 
     Object columnValue = handleSpannerColumnType(spannerType, columnName, valuesJson);
 
-    if (columnValue == null) {
-      LOG.warn("Column value is null for column: {}, type: {}", columnName, spannerType);
-      return PreparedStatementValueObject.create(cassandraType, NullClass.INSTANCE);
-    }
-    return PreparedStatementValueObject.create(cassandraType, columnValue);
+    return PreparedStatementValueObject.create(
+        cassandraType, Objects.requireNonNullElse(columnValue, NullClass.INSTANCE));
   }
 
   /**
@@ -673,7 +673,9 @@ public static PreparedStatementValueObject<?> getColumnValueByType(
    */
   public static Object castToExpectedType(String cassandraType, Object columnValue) {
     try {
-      return parseAndCastToCassandraType(cassandraType, columnValue).value();
+      PreparedStatementValueObject<?> valueObject =
+          parseAndCastToCassandraType(cassandraType, columnValue);
+      return valueObject != null ? valueObject.value() : null;
     } catch (IllegalArgumentException e) {
       LOG.error("Error converting value for column: {}, type: {}", cassandraType, e.getMessage());
       throw new IllegalArgumentException(

v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/SpannerToCassandraSourceDBCustomTransformationIT.java

@@ -30,6 +30,7 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.HashSet;
+import java.util.Objects;
 import org.apache.beam.it.cassandra.CassandraResourceManager;
 import org.apache.beam.it.common.PipelineLauncher;
 import org.apache.beam.it.common.PipelineOperator;
@@ -202,6 +203,8 @@ private void assertBasicRowInCassandraDB() throws InterruptedException {
       assertThat(row.getString("full_name")).isEqualTo("Jone Woe");
       assertThat(row.getString("first_name")).isEqualTo("Jone");
       assertThat(row.getString("last_name")).isEqualTo("Woe");
+      assertThat(Objects.requireNonNull(row.getString("empty_string")).isEmpty()).isTrue();
+      assertThat(row.isNull("null_key")).isTrue();
       assertThat(row.getInt("id")).isEqualTo(1);
     }
   }