Remove special handling for only PK column tables (#2929)

* Remove special handling for only PK column tables

* Fix spotless

* 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,44 +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 = "";
 
-    for (Map.Entry<String, String> entry : pkcolumnNameValues.entrySet()) {
-      String colName = entry.getKey();
-      String colValue = entry.getValue();
-
-      allColumns += "`" + colName + "`,";
-      allValues += colValue + ",";
-    }
-
-    if (columnNameValues.size() == 0) { // if there are only PKs
-      // trim the last ','
-      allColumns = allColumns.substring(0, allColumns.length() - 1);
-      allValues = allValues.substring(0, allValues.length() - 1);
-
-      String returnVal =
-          "INSERT INTO `" + tableName + "`(" + allColumns + ")" + " VALUES (" + allValues + ") ";
-      return new DMLGeneratorResponse(returnVal);
-    }
     int 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;
-      if (!primaryKeys.contains(colName)) {
-        updateValues += " `" + colName + "` = " + colValue;
-      }
+      updateValues += " `" + colName + "` = " + colValue;
 
-      if (index + 1 < columnNameValues.size()) {
+      // Add comma if not the last item in this loop
+      if (index + 1 < allColumnNameValues.size()) {
         allColumns += ",";
         allValues += ",";
         updateValues += ",";
@@ -214,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;
@@ -256,6 +258,42 @@ public void primaryKeyNotPresentInSourceSchema() {
     assertTrue(sql.isEmpty());
   }
 
+  @Test
+  public void tableOnlyContainsPrimaryKeyColumns() {
+    String sessionFile = "src/test/resources/onlyPKColumnsSession.json";
+    Ddl ddl = SchemaUtils.buildSpannerDdlFromSessionFile(sessionFile);
+    SourceSchema sourceSchema = SchemaUtils.buildSourceSchemaFromSessionFile(sessionFile);
+    ISchemaMapper schemaMapper = new SessionBasedMapper(sessionFile, ddl);
+
+    String tableName = "resource_access";
+    String newValuesString = "{\"user_id\":\"101\",\"group_id\":\"5\",\"resource_id\":\"99\"}";
+    JSONObject newValuesJson = new JSONObject(newValuesString);
+    // the keys and the newValues are the same because all the columns are part of the key
+    JSONObject keyValuesJson = new JSONObject(newValuesString);
+    String modType = "INSERT";
+
+    /*The expected sql is:
+    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
+    */
+    MySQLDMLGenerator mySQLDMLGenerator = new MySQLDMLGenerator();
+    DMLGeneratorResponse dmlGeneratorResponse =
+        mySQLDMLGenerator.getDMLStatement(
+            new DMLGeneratorRequest.Builder(
+                    modType, tableName, newValuesJson, keyValuesJson, "+00:00")
+                .setSchemaMapper(schemaMapper)
+                .setDdl(ddl)
+                .setSourceSchema(sourceSchema)
+                .build());
+    String sql = dmlGeneratorResponse.getDmlStatement();
+    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
   public void timezoneOffsetMismatch() {
     String sessionFile = "src/test/resources/timeZoneSession.json";
@@ -965,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
@@ -994,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
@@ -1249,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
@@ -1264,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();
+  }
 }