The PreparedStatement.setObject method supports the BigInteger type.

jd-opensource · Aug 11, 2023 · fb358a6 · fb358a6
1 parent 68f0db6
commit fb358a6
Show file tree

Hide file tree

Showing 11 changed files with 138 additions and 43 deletions.
diff --git a/src/main/java/com/jd/jdbc/Executor.java b/src/main/java/com/jd/jdbc/Executor.java
@@ -33,6 +33,7 @@
 import com.jd.jdbc.queryservice.StreamIterator;
 import com.jd.jdbc.queryservice.util.RoleUtils;
 import com.jd.jdbc.session.SafeSession;
+import com.jd.jdbc.session.VitessSession;
 import com.jd.jdbc.sqlparser.Comment;
 import com.jd.jdbc.sqlparser.SQLUtils;
 import com.jd.jdbc.sqlparser.SqlParser;
@@ -66,6 +67,7 @@
 import io.vitess.proto.Query;
 import io.vitess.proto.Topodata;
 import io.vitess.proto.Vtgate;
+import java.math.BigInteger;
 import java.sql.SQLException;
 import java.sql.SQLFeatureNotSupportedException;
 import java.util.ArrayList;
@@ -404,9 +406,11 @@ private void saveSessionStats(SafeSession safeSession, VtSqlStatementType stmtTy
         if (e != null) {
             return;
         }
+
+        VitessSession sessionNew = safeSession.getVitessConnection().getSessionNew();
         Vtgate.Session.Builder sessionBuilder = safeSession.getVitessConnection().getSession().toBuilder().setFoundRows(resultSet.getRowsAffected());
-        if (resultSet.getInsertID() > 0) {
-            sessionBuilder.setLastInsertId(resultSet.getInsertID());
+        if (resultSet.getInsertID().compareTo(BigInteger.ZERO) > 0) {
+            sessionNew.setLastInsertId(resultSet.getInsertID());
         }
         switch (stmtType) {
             case StmtInsert:
@@ -425,6 +429,7 @@ private void saveSessionStats(SafeSession safeSession, VtSqlStatementType stmtTy
             default:
                 break;
         }
+        safeSession.getVitessConnection().setSessionNew(sessionNew);
         safeSession.getVitessConnection().setSession(sessionBuilder.build());
     }
 

diff --git a/src/main/java/com/jd/jdbc/queryservice/NativeQueryService.java b/src/main/java/com/jd/jdbc/queryservice/NativeQueryService.java
@@ -43,6 +43,7 @@
 import io.prometheus.client.Histogram;
 import io.vitess.proto.Query;
 import io.vitess.proto.Topodata;
+import java.math.BigInteger;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
@@ -490,7 +491,7 @@ private VtResultSet toVtResultSet(final boolean isQuery, final Statement stateme
             ret.setRowsAffected(statement.getUpdateCount());
             ResultSet resultSet = statement.getGeneratedKeys();
             if (resultSet.next()) {
-                ret.setInsertID(resultSet.getLong(1));
+                ret.setInsertID(resultSet.getObject(1, BigInteger.class));
             }
         }
         return ret;

diff --git a/src/main/java/com/jd/jdbc/session/SafeSession.java b/src/main/java/com/jd/jdbc/session/SafeSession.java
@@ -64,6 +64,7 @@ public static SafeSession newSafeSession(VitessConnection vitessConnection, Vtga
         }
         SafeSession safeSession = new SafeSession();
         safeSession.setVitessConnection(new VitessConnection(vitessConnection.getResolver(), sessn));
+        safeSession.vitessConnection.setSessionNew(vitessConnection.getSessionNew());
         return safeSession;
     }
 
@@ -85,6 +86,7 @@ public static SafeSession newAutoCommitSession(VitessConnection vitessConnection
     public void resetTx() {
         this.lock.lock();
         try {
+            VitessSession currentSessionNew = this.vitessConnection.getSessionNew();
             Vtgate.Session.Builder resetTxSessionBuilder = this.vitessConnection.getSession().toBuilder();
             this.mustRollback = false;
             this.autocommitState = AutocommitState.NOT_AUTO_COMMITTABLE;
@@ -97,6 +99,7 @@ public void resetTx() {
                 resetTxSessionBuilder.clearPostSessions();
             }
             this.vitessConnection.setSession(resetTxSessionBuilder.build());
+            this.vitessConnection.setSessionNew(currentSessionNew);
         } finally {
             this.lock.unlock();
         }
@@ -108,6 +111,7 @@ public void resetTx() {
     public void reset() {
         this.lock.lock();
         try {
+            VitessSession currentSessionNew = this.vitessConnection.getSessionNew();
             Vtgate.Session.Builder resetSessionBuilder = this.vitessConnection.getSession().toBuilder();
             this.mustRollback = false;
             this.autocommitState = AutocommitState.NOT_AUTO_COMMITTABLE;
@@ -118,6 +122,7 @@ public void reset() {
             resetSessionBuilder.clearPreSessions();
             resetSessionBuilder.clearPostSessions();
             this.vitessConnection.setSession(resetSessionBuilder.build());
+            this.vitessConnection.setSessionNew(currentSessionNew);
         } finally {
             this.lock.unlock();
         }
@@ -267,6 +272,7 @@ public void appendOrUpdate(Vtgate.Session.ShardSession shardSession, Vtgate.Tran
                 case NORMAL:
                     shardSessionList = buildSession(shardSession, this.vitessConnection.getSession().getShardSessionsList());
                     this.vitessConnection.setSession(this.vitessConnection.getSession().toBuilder().clearShardSessions().addAllShardSessions(shardSessionList).build());
+                    this.vitessConnection.setSessionNew(this.vitessConnection.getSessionNew());
 
                     // isSingle is enforced only for normmal commit order operations.
                     if (this.isSingleDb(txMode) && this.vitessConnection.getSession().getShardSessionsCount() > 1) {
@@ -277,10 +283,12 @@ public void appendOrUpdate(Vtgate.Session.ShardSession shardSession, Vtgate.Tran
                 case PRE:
                     shardSessionList = buildSession(shardSession, this.vitessConnection.getSession().getPreSessionsList());
                     this.vitessConnection.setSession(this.vitessConnection.getSession().toBuilder().clearPreSessions().addAllPreSessions(shardSessionList).build());
+                    this.vitessConnection.setSessionNew(this.vitessConnection.getSessionNew());
                     break;
                 case POST:
                     shardSessionList = buildSession(shardSession, this.vitessConnection.getSession().getPostSessionsList());
                     this.vitessConnection.setSession(this.vitessConnection.getSession().toBuilder().clearPostSessions().addAllPostSessions(shardSessionList).build());
+                    this.vitessConnection.setSessionNew(this.vitessConnection.getSessionNew());
                     break;
                 default:
                     throw new SQLException("BUG: SafeSession.AppendOrUpdate: unexpected commitOrder");

diff --git a/src/main/java/com/jd/jdbc/session/VitessSession.java b/src/main/java/com/jd/jdbc/session/VitessSession.java
@@ -0,0 +1,41 @@
+/*
+Copyright 2021 JD Project Authors. Licensed under Apache-2.0.
+
+Copyright 2019 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package com.jd.jdbc.session;
+
+import java.math.BigInteger;
+import lombok.Getter;
+
+public class VitessSession {
+
+    @Getter
+    private BigInteger lastInsertId;
+
+    public VitessSession() {
+        this.lastInsertId = BigInteger.ZERO;
+    }
+
+    public void setLastInsertId(long setId) {
+        this.lastInsertId = BigInteger.valueOf(setId);
+    }
+
+    public void setLastInsertId(BigInteger setId) {
+        this.lastInsertId = setId;
+    }
+
+}
diff --git a/src/main/java/com/jd/jdbc/sqltypes/VtResultSet.java b/src/main/java/com/jd/jdbc/sqltypes/VtResultSet.java
@@ -20,6 +20,7 @@
 
 import com.jd.jdbc.vitess.VitessStatement;
 import io.vitess.proto.Query;
+import java.math.BigInteger;
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -34,7 +35,7 @@ public class VtResultSet implements VtRowList {
 
     private long rowsAffected;
 
-    private long insertID;
+    private BigInteger insertID;
 
     private List<List<VtResultValue>> rows;
 
@@ -48,21 +49,24 @@ public class VtResultSet implements VtRowList {
 
     public VtResultSet() {
         this.rows = new ArrayList<>();
+        this.insertID = BigInteger.valueOf(0);
     }
 
-    public VtResultSet(long insertID, long rowsAffected) {
+    public VtResultSet(BigInteger insertID, long rowsAffected) {
         this.insertID = insertID;
         this.rowsAffected = rowsAffected;
     }
 
     public VtResultSet(Query.Field[] fields, List<List<VtResultValue>> rows) {
         this.fields = fields;
         this.rows = rows;
+        this.insertID = BigInteger.valueOf(0);
     }
 
     public VtResultSet(long rowsAffected, List<List<VtResultValue>> rows) {
         this.rowsAffected = rowsAffected;
         this.rows = rows;
+        this.insertID = BigInteger.valueOf(0);
     }
 
     @Override
@@ -107,7 +111,7 @@ private void appendResult(VtResultSet src, boolean ignoreTable) throws SQLExcept
             }
         }
         this.rowsAffected += src.rowsAffected;
-        if (src.insertID != 0) {
+        if (!src.insertID.equals(BigInteger.valueOf(0))) {
             this.insertID = src.insertID;
         }
         if (this.rows == null || this.rows.isEmpty()) {
@@ -227,4 +231,17 @@ public VtRowList clone() {
         VtRowList vtRowList = new VtResultSet(this.fields, this.rows);
         return vtRowList;
     }
+
+    public BigInteger getInsertID() {
+        return this.insertID;
+    }
+
+    public void setInsertID(long setId) {
+        this.insertID = BigInteger.valueOf(setId);
+    }
+
+    public void setInsertID(BigInteger setId) {
+        this.insertID = setId;
+    }
+
 }
diff --git a/src/main/java/com/jd/jdbc/sqltypes/VtRowList.java b/src/main/java/com/jd/jdbc/sqltypes/VtRowList.java
@@ -19,6 +19,7 @@
 package com.jd.jdbc.sqltypes;
 
 import io.vitess.proto.Query;
+import java.math.BigInteger;
 import java.sql.SQLException;
 import java.util.List;
 
@@ -28,7 +29,7 @@ public interface VtRowList extends VtResultSetMessage {
 
     long getRowsAffected();
 
-    long getInsertID();
+    BigInteger getInsertID();
 
     boolean isQuery();
 

diff --git a/src/main/java/com/jd/jdbc/sqltypes/VtStreamResultSet.java b/src/main/java/com/jd/jdbc/sqltypes/VtStreamResultSet.java
@@ -21,6 +21,7 @@
 import com.jd.jdbc.IExecute;
 import com.jd.jdbc.vitess.VitessStatement;
 import io.vitess.proto.Query;
+import java.math.BigInteger;
 import java.sql.SQLException;
 import java.util.List;
 
@@ -69,8 +70,8 @@ public long getRowsAffected() {
     }
 
     @Override
-    public long getInsertID() {
-        return 0;
+    public BigInteger getInsertID() {
+        return BigInteger.valueOf(0);
     }
 
     @Override

diff --git a/src/main/java/com/jd/jdbc/vitess/VitessConnection.java b/src/main/java/com/jd/jdbc/vitess/VitessConnection.java
@@ -29,6 +29,7 @@
 import com.jd.jdbc.queryservice.TabletDialer;
 import com.jd.jdbc.queryservice.util.RoleUtils;
 import com.jd.jdbc.session.SafeSession;
+import com.jd.jdbc.session.VitessSession;
 import com.jd.jdbc.sqlparser.ast.statement.SQLCommitStatement;
 import com.jd.jdbc.sqlparser.ast.statement.SQLRollbackStatement;
 import com.jd.jdbc.sqlparser.support.logging.Log;
@@ -101,6 +102,10 @@ public class VitessConnection extends AbstractVitessConnection {
     @Setter
     private Vtgate.Session session;
 
+    @Getter
+    @Setter
+    private VitessSession sessionNew;
+
     @Getter
     private String defaultKeyspace;
 
@@ -127,6 +132,7 @@ public VitessConnection(String url, Properties prop, TopoServer topoServer, Reso
             .setAutocommit(true)
             .build();
         buildServerSessionPropertiesMap();
+        this.sessionNew = new VitessSession();
         if (log.isDebugEnabled()) {
             log.debug("create VitessConnection");
         }

diff --git a/src/main/java/com/jd/jdbc/vitess/VitessPreparedStatement.java b/src/main/java/com/jd/jdbc/vitess/VitessPreparedStatement.java
@@ -36,6 +36,7 @@
 import java.io.InputStream;
 import java.io.Reader;
 import java.math.BigDecimal;
+import java.math.BigInteger;
 import java.sql.Blob;
 import java.sql.Clob;
 import java.sql.Date;
@@ -314,6 +315,10 @@ public void setObject(int parameterIndex, Object x) throws SQLException {
             setInt(parameterIndex, (Integer) x);
             return;
         }
+        if (x instanceof BigInteger) {
+            setBigDecimal(parameterIndex, new BigDecimal((BigInteger) x));
+            return;
+        }
         if (x instanceof Long) {
             setLong(parameterIndex, (Long) x);
             return;