[CONJ-141] Batch Statement Rewrite: Support for ON DUPLICATE KEY

mariadb-corporation · Jul 22, 2015 · 1bfe353 · 1bfe353
1 parent afa24ed
commit 1bfe353
Show file tree

Hide file tree

Showing 6 changed files with 86 additions and 14 deletions.
diff --git a/src/main/java/org/mariadb/jdbc/MySQLDataSource.java b/src/main/java/org/mariadb/jdbc/MySQLDataSource.java
@@ -224,7 +224,7 @@ public void setProperties(String properties) {
      * Sets the connection string URL.
      *
      * @param url the connection string
-     * @throws SQLException
+     * @throws SQLException if error in URL
      */
     public void setURL(String url) throws SQLException {
         setUrl(url);
@@ -233,7 +233,7 @@ public void setURL(String url) throws SQLException {
     /**
      * Sets the connection string URL.
      * @param s the connection string
-     * @throws SQLException
+     * @throws SQLException if error in URL
      */
     public void setUrl(String s) throws SQLException {
         this.jdbcUrl.parseUrl(s);

diff --git a/src/main/java/org/mariadb/jdbc/internal/common/packet/commands/StreamedQueryPacket.java b/src/main/java/org/mariadb/jdbc/internal/common/packet/commands/StreamedQueryPacket.java
@@ -83,14 +83,18 @@ public int send(final OutputStream ostream) throws IOException, QueryException {
             PacketOutputStream pos = (PacketOutputStream)ostream;
             pos.startPacket(0);
             pos.write(0x03);
-            queries.get(0).writeTo(ostream);
-            for (int i=1;i<queries.size();i++) {
-                if (isRewritable) {
-                    queries.get(i).writeToRewritablePart(ostream, rewriteOffset);
-                } else {
+            if (!isRewritable) {
+                queries.get(0).writeTo(ostream);
+                for (int i=1;i<queries.size();i++) {
                     pos.write(';');
                     queries.get(i).writeTo(ostream);
                 }
+            } else {
+                queries.get(0).writeFirstRewritePart(ostream);
+                for (int i = 1;i<queries.size();i++) {
+                        queries.get(i).writeToRewritablePart(ostream, rewriteOffset);
+                }
+                queries.get(0).writeLastRewritePart(ostream);
             }
             pos.finishPacket();
         }

diff --git a/src/main/java/org/mariadb/jdbc/internal/common/query/MySQLParameterizedQuery.java b/src/main/java/org/mariadb/jdbc/internal/common/query/MySQLParameterizedQuery.java
@@ -67,14 +67,21 @@ public class MySQLParameterizedQuery implements ParameterizedQuery {
     private int paramCount;
     private String query;
     private byte[][] queryPartsArray;
+
     private byte[] rewriteFirstPart = null;
+    private byte[] rewriteRepeatLastPart = null;
+    private byte[] rewriteNotRepeatLastPart = null;
+
 
     public MySQLParameterizedQuery(String query, boolean noBackslashEscapes, int rewriteOffset) {
         this.query = query;
         List<String> queryParts = createQueryParts(query, noBackslashEscapes);
         if (rewriteOffset != -1) {
             try {
                 rewriteFirstPart = queryParts.get(0).substring(rewriteOffset + 1).getBytes("UTF-8");
+                String lastPart = queryParts.get(queryParts.size() - 1 );
+                rewriteRepeatLastPart = lastPart.substring(0, lastPart.indexOf(")")).getBytes("UTF-8");
+                rewriteNotRepeatLastPart = lastPart.substring(lastPart.indexOf(")") + 1 ).getBytes("UTF-8");
             } catch (UnsupportedEncodingException e) {
                 throw new RuntimeException("UTF-8 not supported", e);
             }
@@ -104,6 +111,8 @@ public MySQLParameterizedQuery cloneQuery() {
         q.query = query;
         q.queryPartsArray = queryPartsArray;
         q.rewriteFirstPart = rewriteFirstPart;
+        q.rewriteRepeatLastPart = rewriteRepeatLastPart;
+        q.rewriteNotRepeatLastPart = rewriteNotRepeatLastPart;
         return q;
     }
 
@@ -142,18 +151,39 @@ public void writeTo(final OutputStream os) throws IOException, QueryException {
         }
     }
 
+    public void writeFirstRewritePart(final OutputStream os) throws IOException, QueryException {
+        if(queryPartsArray.length == 0) {
+            throw new AssertionError("Invalid query, queryParts was empty");
+        }
+
+        for(int i = 0; i<queryPartsArray.length - 1; i++) {
+            os.write(queryPartsArray[i]);
+            parameters[i].writeTo(os);
+        }
+        if (rewriteRepeatLastPart != null) os.write(rewriteRepeatLastPart);
+        os.write(41); // ")" in UTF-8
+    }
+
+    public void writeLastRewritePart(final OutputStream os) throws IOException, QueryException {
+        if(rewriteNotRepeatLastPart != null) {
+            os.write(rewriteNotRepeatLastPart);
+        }
+    }
 
     public void writeToRewritablePart(final OutputStream os, int rewriteOffset) throws IOException, QueryException {
         if(queryPartsArray.length == 0) {
             throw new AssertionError("Invalid query, queryParts was empty");
         }
-        os.write((",(").getBytes());
+
+        os.write(new byte[]{44,40}); //",(" in UTF-8
         os.write(rewriteFirstPart);
-        for(int i = 1; i<queryPartsArray.length; i++) {
-            parameters[i-1].writeTo(os);
-            if(queryPartsArray[i].length != 0)
-                os.write(queryPartsArray[i]);
+        for(int i = 0; i < parameters.length ; i++) {
+            parameters[i].writeTo(os);
+            if (i < parameters.length - 1)
+                os.write(queryPartsArray[i + 1]);
+            else os.write(rewriteRepeatLastPart);
         }
+        os.write(41); // ")" in UTF-8
     }
 
     private boolean containsNull(ParameterHolder[] parameters) {

diff --git a/src/main/java/org/mariadb/jdbc/internal/common/query/MySQLQuery.java b/src/main/java/org/mariadb/jdbc/internal/common/query/MySQLQuery.java
@@ -72,6 +72,13 @@ public void writeTo(final OutputStream os) throws IOException {
         }
     }
 
+    public void writeFirstRewritePart(final OutputStream os) throws IOException, QueryException {
+        writeTo(os);
+    }
+
+    public void writeLastRewritePart(final OutputStream os) throws IOException, QueryException { }
+
+
     public void writeToRewritablePart(final OutputStream os, int rewriteOffset) throws IOException, QueryException {
         try {
             byte[] queryToSend = query.substring(rewriteOffset).getBytes("UTF-8");
@@ -102,10 +109,9 @@ public void writeTo(OutputStream ostream, int offset, int packLength) throws IOE
         } catch (UnsupportedEncodingException e) {
             throw new RuntimeException("Unsupported encoding: " + e.getMessage(), e);
         }
-
-
     }
 
+
     public void validate() throws QueryException{
 
     }

diff --git a/src/main/java/org/mariadb/jdbc/internal/common/query/Query.java b/src/main/java/org/mariadb/jdbc/internal/common/query/Query.java
@@ -56,6 +56,8 @@ WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWIS
 public interface Query {
     String getQuery();
     void writeTo(OutputStream os) throws IOException, QueryException;
+    void writeFirstRewritePart(final OutputStream os) throws IOException, QueryException;
+    void writeLastRewritePart(final OutputStream os) throws IOException, QueryException;
     void writeToRewritablePart(OutputStream os, int rewriteOffset) throws IOException, QueryException;
     QueryType getQueryType();
     void validate() throws QueryException;

diff --git a/src/test/java/org/mariadb/jdbc/MultiTest.java b/src/test/java/org/mariadb/jdbc/MultiTest.java
@@ -238,6 +238,36 @@ private int retrieveSessionVariableFromServer(Connection tmpConnection, String v
         throw new RuntimeException("Unable to retrieve, variable value from Server " + variable);
     }
 
+
+
+    @Test
+    public void rewriteBatchedStatementsWithQueryFirstAndLAst() throws SQLException {
+        Properties props = new Properties();
+        props.setProperty("rewriteBatchedStatements", "true");
+        Connection tmpConnection = null;
+        try {
+            tmpConnection = openNewConnection(connURI, props);
+            Statement st = tmpConnection.createStatement();
+            st.executeUpdate("drop table if exists t3_dupp");
+            st.executeUpdate("create table t3_dupp(col1 int, pkey int NOT NULL, col2 int, col3 int, col4 int, PRIMARY KEY (`pkey`))");
+
+            PreparedStatement sqlInsert = connection.prepareStatement("INSERT INTO t3_dupp(col1, pkey,col2,col3,col4) VALUES (9, ?, 5, ?, 8) ON DUPLICATE KEY UPDATE pkey=pkey+10");
+            sqlInsert.setInt(1, 1);
+            sqlInsert.setInt(2, 2);
+            sqlInsert.addBatch();
+
+            sqlInsert.setInt(1, 2);
+            sqlInsert.setInt(2, 5);
+            sqlInsert.addBatch();
+
+            sqlInsert.setInt(1, 7);
+            sqlInsert.setInt(2, 6);
+            sqlInsert.addBatch();
+            sqlInsert.executeBatch();
+        } finally {
+            if (tmpConnection != null) tmpConnection.close();
+        }
+    }
     /**
      * CONJ-142: Using a semicolon in a string with "rewriteBatchedStatements=true" fails
      *