Refactor maps and update UTs

Author: manitgupta

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

@@ -121,53 +121,23 @@ public DMLGeneratorResponse getDMLStatement(DMLGeneratorRequest dmlGeneratorRequ
   }
 
   private static DMLGeneratorResponse getUpsertStatement(
-      String tableName,
-      List<String> primaryKeys,
-      Map<String, String> columnNameValues,
-      Map<String, String> pkcolumnNameValues) {
+      String tableName, Map<String, String> allColumnNameValues) {
 
     String allColumns = "";
     String allValues = "";
     String updateValues = "";
 
     int index = 0;
 
-    for (Map.Entry<String, String> entry : pkcolumnNameValues.entrySet()) {
-      String colName = entry.getKey();
-      String colValue = entry.getValue();
-
-      allColumns += "`" + colName + "`";
-      allValues += colValue;
-      // Add PK columns as well to the UPDATE clause
-      updateValues += " `" + colName + "` = " + colValue;
-
-      // Add comma if not the last item in this loop
-      if (index + 1 < pkcolumnNameValues.size()) {
-        allColumns += ",";
-        allValues += ",";
-        updateValues += ",";
-      }
-      index++;
-    }
-
-    // Add a separator comma if we had PKs AND we have non-PKs to add
-    if (!pkcolumnNameValues.isEmpty() && !columnNameValues.isEmpty()) {
-      allColumns += ",";
-      allValues += ",";
-      updateValues += ",";
-    }
-
-    index = 0;
-
-    for (Map.Entry<String, String> entry : columnNameValues.entrySet()) {
+    for (Map.Entry<String, String> entry : allColumnNameValues.entrySet()) {
       String colName = entry.getKey();
       String colValue = entry.getValue();
       allColumns += "`" + colName + "`";
       allValues += colValue;
       updateValues += " `" + colName + "` = " + colValue;
 
       // Add comma if not the last item in this loop
-      if (index + 1 < columnNameValues.size()) {
+      if (index + 1 < allColumnNameValues.size()) {
         allColumns += ",";
         allValues += ",";
         updateValues += ",";
@@ -223,8 +193,8 @@ private static DMLGeneratorResponse generateUpsertStatement(
             dmlGeneratorRequest.getKeyValuesJson(),
             dmlGeneratorRequest.getSourceDbTimezoneOffset(),
             dmlGeneratorRequest.getCustomTransformationResponse());
-    return getUpsertStatement(
-        sourceTable.name(), sourceTable.primaryKeyColumns(), columnNameValues, pkcolumnNameValues);
+    columnNameValues.putAll(pkcolumnNameValues);
+    return getUpsertStatement(sourceTable.name(), columnNameValues);
   }
 
   private static Map<String, String> getColumnValues(

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

@@ -16,6 +16,7 @@
 package com.google.cloud.teleport.v2.templates.dbutils.dml;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
@@ -34,6 +35,7 @@
 import java.io.InputStream;
 import java.nio.channels.Channels;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.beam.sdk.io.FileSystems;
@@ -283,9 +285,13 @@ public void tableOnlyContainsPrimaryKeyColumns() {
                 .setSourceSchema(sourceSchema)
                 .build());
     String sql = dmlGeneratorResponse.getDmlStatement();
-    assertTrue(
-        sql.contains(
-            "INSERT INTO `resource_access`(`user_id`,`group_id`,`resource_id`) VALUES (101,5,99) ON DUPLICATE KEY UPDATE  `user_id` = 101, `group_id` = 5, `resource_id` = 99"));
+    assertThat(sql.contains("ON DUPLICATE KEY UPDATE"));
+    assertTrue(sql.contains("`user_id` = 101"));
+    assertTrue(sql.contains("`group_id` = 5"));
+    assertTrue(sql.contains("`resource_id` = 99"));
+    assertEquals(2, countInSQL(sql, "user_id"));
+    assertEquals(2, countInSQL(sql, "group_id"));
+    assertEquals(2, countInSQL(sql, "resource_id"));
   }
 
   @Test
@@ -997,10 +1003,13 @@ public void testSpannerKeyIsNull() {
                 .setSourceSchema(sourceSchema)
                 .build());
     String sql = dmlGeneratorResponse.getDmlStatement();
-
-    assertTrue(
-        sql.contains(
-            "INSERT INTO `Singers`(`SingerId`,`FirstName`,`LastName`) VALUES (NULL,'kk','ll')"));
+    assertThat(sql.contains("ON DUPLICATE KEY UPDATE"));
+    assertTrue(sql.contains("`FirstName` = 'kk'"));
+    assertTrue(sql.contains("`SingerId` = NULL"));
+    assertTrue(sql.contains("`LastName` = 'll'"));
+    assertEquals(2, countInSQL(sql, "FirstName"));
+    assertEquals(2, countInSQL(sql, "SingerId"));
+    assertEquals(2, countInSQL(sql, "LastName"));
   }
 
   @Test
@@ -1026,9 +1035,13 @@ public void testKeyInNewValuesJson() {
                 .setSourceSchema(sourceSchema)
                 .build());
     String sql = dmlGeneratorResponse.getDmlStatement();
-    assertTrue(
-        sql.contains(
-            "INSERT INTO `Singers`(`SingerId`,`FirstName`,`LastName`) VALUES (NULL,'kk','ll')"));
+    assertThat(sql.contains("ON DUPLICATE KEY UPDATE"));
+    assertTrue(sql.contains("`FirstName` = 'kk'"));
+    assertTrue(sql.contains("`SingerId` = NULL"));
+    assertTrue(sql.contains("`LastName` = 'll'"));
+    assertEquals(2, countInSQL(sql, "FirstName"));
+    assertEquals(2, countInSQL(sql, "SingerId"));
+    assertEquals(2, countInSQL(sql, "LastName"));
   }
 
   @Test
@@ -1281,8 +1294,11 @@ public void customTransformationMatch() {
                 .build());
     String sql = dmlGeneratorResponse.getDmlStatement();
 
+    assertThat(sql.contains("ON DUPLICATE KEY UPDATE"));
     assertTrue(sql.contains("`FullName` = 'kk ll'"));
-    assertTrue(sql.contains("VALUES (1,'kk ll')"));
+    assertTrue(sql.contains("`SingerId` = 1"));
+    assertEquals(2, countInSQL(sql, "FullName"));
+    assertEquals(2, countInSQL(sql, "SingerId"));
   }
 
   @Test
@@ -1296,4 +1312,13 @@ public void testConvertBase64ToXHex() {
         IllegalArgumentException.class,
         () -> MySQLDMLGenerator.convertBase64ToHex("####GOOGLE####"));
   }
+
+  public long countInSQL(String sql, String targetWord) {
+    if (sql == null || sql.isEmpty()) {
+      return 0;
+    }
+    return Arrays.stream(sql.split("\\W+"))
+        .filter(word -> word.equalsIgnoreCase(targetWord))
+        .count();
+  }
 }