Factor sql generation out of BinCopyUploader

Author: joerivanruth

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopyDataWriterFactory.java

@@ -34,21 +34,22 @@ public class BinCopyDataWriterFactory implements DataWriterFactory, Serializable
 	private static final long serialVersionUID = 0L;
 
 	private final Parms parms;
-	private final String[] columns;
 	private final Step[] steps;
+	private final BinCopySql sqlstmt;
 
-	public BinCopyDataWriterFactory(Parms parms, String[] columns, Step[] steps) {
+	public BinCopyDataWriterFactory(Parms parms, Step[] steps, BinCopySql sqlstmt) {
 		this.parms = parms;
-		this.columns = columns;
 		this.steps = steps;
+		this.sqlstmt = sqlstmt;
 	}
 
 	@Override
 	public DataWriter<InternalRow> createWriter(int partitionId, long taskId) {
 		try {
 			String identifier = "part" + partitionId + "-task" + taskId;
 			Collector collector = new Collector();
-			BinCopyUploader uploader = new BinCopyUploader(parms.getDestination(), collector, identifier, columns);
+			sqlstmt.identifier(identifier);
+			BinCopyUploader uploader = new BinCopyUploader(parms.getDestination(), collector, sqlstmt);
 			collector.registerWithConverters(steps);
 			return new MonetDataWriter(collector, steps, uploader, parms.isImmediateCommit(), identifier, parms.getBatchSize());
 		} catch (SQLException e) {

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopySql.java

@@ -0,0 +1,88 @@
+/*
+ * SPDX-License-Identifier: MPL-2.0
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0.  If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Copyright MonetDB Solutions B.V.
+ */
+
+package org.monetdb.spark.bincopy;
+
+import java.io.PrintWriter;
+import java.io.Serializable;
+import java.io.StringWriter;
+
+public class BinCopySql implements Serializable {
+	String quotedTableName;
+	String[] columnNames;
+	String identifier;
+	String prefix = "";
+	String suffix = "";
+	boolean onClient = true;
+
+	public BinCopySql(String quotedTableName, String[] columnNames) {
+		this.quotedTableName = quotedTableName;
+		this.columnNames = columnNames.clone();
+	}
+
+	public BinCopySql identifier(String identifier) {
+		this.identifier = identifier;
+		return this;
+	}
+
+	public BinCopySql suffix(String suffix) {
+		this.suffix = suffix;
+		return this;
+	}
+
+	public BinCopySql prefix(String prefix) {
+		this.prefix = prefix;
+		return this;
+	}
+
+	public BinCopySql onClient(boolean onClient) {
+		this.onClient = onClient;
+		return this;
+	}
+
+	public BinCopySql onServer(boolean onServer) {
+		return onClient(!onServer);
+	}
+
+	@Override
+	public String toString() {
+		String sep;
+
+		StringWriter sw = new StringWriter();
+		PrintWriter pw = new PrintWriter(sw);
+
+		String identifierComment = identifier != null ? " /* " + identifier.replace("*/", "* /") + " */" : "";
+		pw.printf("COPY%s LITTLE ENDIAN BINARY INTO %s (", identifierComment, quotedTableName);
+		sep = "";
+		for (String col : columnNames) {
+			pw.printf("%s\"%s\"", sep, col.replace("\"", "\"\""));
+			sep = ", ";
+		}
+		pw.println(")");
+
+		pw.print("FROM ");
+		final String qprefix, qsuffix;
+		if (prefix.contains("\\") || prefix.contains("'") || suffix.contains("\\") || suffix.contains("'")) {
+			qprefix = "R'" + prefix.replace("'", "''");
+			qsuffix = suffix.replace("'", "''") + "'";
+		} else {
+			qprefix = "'" + prefix;
+			qsuffix = suffix + "'";
+		}
+		sep = "";
+		for (int i = 0; i < columnNames.length; i++) {
+			pw.printf("%s%s%d%s", sep, qprefix, i, qsuffix);
+			sep = ", ";
+		} pw.println();
+		pw.println(onClient ? "ON CLIENT" : "ON SERVER");
+
+		return sw.toString();
+	}
+}

monetdb-spark/src/main/java/org/monetdb/spark/bincopy/BinCopyUploader.java

@@ -14,8 +14,6 @@
 import org.monetdb.spark.common.Destination;
 import org.monetdb.spark.workerside.Collector;
 
-import java.io.PrintWriter;
-import java.io.StringWriter;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
 
@@ -24,34 +22,14 @@ public class BinCopyUploader {
 	private final MonetConnection conn;
 	private final PreparedStatement stmt;
 
-	public BinCopyUploader(Destination dest, Collector collector, String identifier, String[] columns) throws SQLException {
+	public BinCopyUploader(Destination dest, Collector collector, BinCopySql sqlstmt) throws SQLException {
 		String sep;
 		this.collector = collector;
 		this.conn = dest.connect();
 		conn.setAutoCommit(false);
 		conn.setUploadHandler(collector);
 
-		StringWriter sw = new StringWriter();
-		PrintWriter pw = new PrintWriter(sw);
-
-		pw.printf("COPY /* %s */ LITTLE ENDIAN BINARY INTO %s (", identifier, dest.getTable());
-		sep = "";
-		for (String col : columns) {
-			pw.printf("%s\"%s\"", sep, col.replace("\"", "\"\""));
-			sep = ", ";
-		}
-		pw.println(")");
-
-		pw.print("FROM ");
-		sep = "";
-		for (int i = 0; i < columns.length; i++) {
-			pw.printf("%s'%d'", sep, i);
-			sep = ", ";
-		}
-		pw.println();
-		pw.println("ON CLIENT");
-
-		String sql = sw.toString();
+		String sql = sqlstmt.toString();
 		stmt = conn.prepareStatement(sql);
 	}
 

monetdb-spark/src/main/java/org/monetdb/spark/driverside/MonetBatchWrite.java

@@ -9,6 +9,7 @@
 import org.apache.spark.sql.connector.write.PhysicalWriteInfo;
 import org.apache.spark.sql.connector.write.WriterCommitMessage;
 import org.monetdb.spark.bincopy.BinCopyDataWriterFactory;
+import org.monetdb.spark.bincopy.BinCopySql;
 import org.monetdb.spark.workerside.Step;
 
 /**
@@ -23,18 +24,18 @@
  */
 public class MonetBatchWrite implements BatchWrite {
 	private final Parms parms;
-	private final String[] columnNames;
 	private final Step[] steps;
+	private final BinCopySql sqlstmt;
 
-	public MonetBatchWrite(Parms parms, String[] columnNames, Step[] steps) {
+	public MonetBatchWrite(Parms parms, Step[] steps, BinCopySql sqlstmt) {
 		this.parms = parms;
-		this.columnNames = columnNames;
 		this.steps = steps;
+		this.sqlstmt = sqlstmt;
 	}
 
 	@Override
 	public DataWriterFactory createBatchWriterFactory(PhysicalWriteInfo info) {
-		return new BinCopyDataWriterFactory(parms, columnNames, steps);
+		return new BinCopyDataWriterFactory(parms, steps, sqlstmt);
 	}
 
 	@Override

monetdb-spark/src/main/java/org/monetdb/spark/driverside/MonetWrite.java

@@ -7,6 +7,7 @@
 import org.apache.spark.sql.connector.metric.CustomMetric;
 import org.apache.spark.sql.connector.write.BatchWrite;
 import org.apache.spark.sql.connector.write.Write;
+import org.monetdb.spark.bincopy.BinCopySql;
 import org.monetdb.spark.bincopy.PlanBuilder;
 import org.monetdb.spark.common.ColumnDescr;
 import org.monetdb.spark.common.Destination;
@@ -53,7 +54,8 @@ public MonetWrite(Parms parms) {
 
 	@Override
 	public BatchWrite toBatch() {
-		return new MonetBatchWrite(parms, builder.getColumns(), builder.getPlan());
+		BinCopySql sqlstmt = new BinCopySql(this.parms.getDestination().getTable(), builder.getColumns());
+		return new MonetBatchWrite(parms, builder.getPlan(), sqlstmt);
 	}
 
 	@Override

monetdb-spark/src/test/java/org/monetdb/spark/bincopy/UploadTests.java

@@ -65,7 +65,8 @@ public void testUploadingData() throws SQLException, ConversionError, IOExceptio
 		Collector collector = new Collector();
 		collector.registerWithConverters(steps);
 		String identifier = "uploadtest";
-		BinCopyUploader uploader = new BinCopyUploader(dest, collector, identifier, builder.getColumns());
+		BinCopySql sqlstmt = new BinCopySql(dest.getTable(), builder.getColumns()).identifier(identifier);
+		BinCopyUploader uploader = new BinCopyUploader(dest, collector, sqlstmt);
 		long batchSize = Long.MAX_VALUE;
 		MonetDataWriter dataWriter = new MonetDataWriter(collector, steps, uploader, false, identifier, batchSize);