Escape column names

cloudsufi · Oct 1, 2024 · fbd7ba6 · fbd7ba6
1 parent 9641b17
commit fbd7ba6
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLSink.java b/cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLSink.java
@@ -25,6 +25,7 @@
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.annotation.Plugin;
 import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
 import io.cdap.cdap.etl.api.batch.BatchSink;
@@ -40,7 +41,11 @@
 import io.cdap.plugin.util.CloudSQLUtil;
 import io.cdap.plugin.util.DBUtils;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
+import java.util.StringJoiner;
 import javax.annotation.Nullable;
 
 /** Sink support for a CloudSQL MySQL database. */
@@ -52,6 +57,7 @@
 public class CloudSQLMySQLSink extends AbstractDBSink<CloudSQLMySQLSink.CloudSQLMySQLSinkConfig> {
 
   private final CloudSQLMySQLSinkConfig cloudsqlMysqlSinkConfig;
+  private static final Character ESCAPE_CHAR = '`';
 
   public CloudSQLMySQLSink(CloudSQLMySQLSinkConfig cloudsqlMysqlSinkConfig) {
     super(cloudsqlMysqlSinkConfig);
@@ -78,6 +84,19 @@ protected DBRecord getDBRecord(StructuredRecord output) {
     return new MysqlDBRecord(output, columnTypes);
   }
 
+  @Override
+  protected void setColumnsInfo(List<Schema.Field> fields) {
+    List<String> columnsList = new ArrayList<>();
+    StringJoiner columnsJoiner = new StringJoiner(",");
+    for (Schema.Field field : fields) {
+      columnsList.add(field.getName());
+      columnsJoiner.add(ESCAPE_CHAR + field.getName() + ESCAPE_CHAR);
+    }
+
+    super.columns = Collections.unmodifiableList(columnsList);
+    super.dbColumns = columnsJoiner.toString();
+  }
+
   @Override
   protected LineageRecorder getLineageRecorder(BatchSinkContext context) {
     String host;