提交 d48c679b authored 作者: Evgenij Ryazanov's avatar Evgenij Ryazanov

Set requested generated keys to client together with result of command

上级 58a30eba
......@@ -14,7 +14,9 @@ import org.h2.engine.Session;
import org.h2.expression.ParameterInterface;
import org.h2.message.DbException;
import org.h2.message.Trace;
import org.h2.result.LocalResult;
import org.h2.result.ResultInterface;
import org.h2.result.ResultWithGeneratedKeys;
import org.h2.util.MathUtils;
/**
......@@ -238,7 +240,7 @@ public abstract class Command implements CommandInterface {
}
@Override
public int executeUpdate(Object generatedKeysRequest) {
public ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest) {
long start = 0;
Database database = session.getDatabase();
Object sync = database.isMultiThreaded() ? (Object) session : (Object) database;
......@@ -257,7 +259,14 @@ public abstract class Command implements CommandInterface {
while (true) {
database.checkPowerOff();
try {
return update();
int updateCount = update();
if (!Boolean.FALSE.equals(generatedKeysRequest)) {
return new ResultWithGeneratedKeys.WithKeys(updateCount,
LocalResult.read(session,
session.getGeneratedKeys().getKeys(),
Integer.MAX_VALUE));
}
return ResultWithGeneratedKeys.of(updateCount);
} catch (DbException e) {
start = filterConcurrentUpdate(e, start);
} catch (OutOfMemoryError e) {
......
......@@ -8,6 +8,7 @@ package org.h2.command;
import java.util.ArrayList;
import org.h2.expression.ParameterInterface;
import org.h2.result.ResultInterface;
import org.h2.result.ResultWithGeneratedKeys;
/**
* Represents a SQL statement.
......@@ -519,7 +520,7 @@ public interface CommandInterface {
*
* @return the update count
*/
int executeUpdate(Object generatedKeysRequest);
ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest);
/**
* Stop the command execution, release all locks and resources
......
......@@ -39,7 +39,7 @@ class CommandList extends Command {
@Override
public int update() {
int updateCount = command.executeUpdate(false);
int updateCount = command.executeUpdate(false).getUpdateCount();
executeRemaining();
return updateCount;
}
......
......@@ -17,6 +17,7 @@ import org.h2.message.DbException;
import org.h2.message.Trace;
import org.h2.result.ResultInterface;
import org.h2.result.ResultRemote;
import org.h2.result.ResultWithGeneratedKeys;
import org.h2.util.New;
import org.h2.value.Transfer;
import org.h2.value.Value;
......@@ -194,10 +195,14 @@ public class CommandRemote implements CommandInterface {
}
@Override
public int executeUpdate(Object generatedKeys) {
public ResultWithGeneratedKeys executeUpdate(Object generatedKeysRequest) {
checkParameters();
boolean supportsGeneratedKeys = session.isSupportsGeneratedKeys();
boolean readGeneratedKeys = supportsGeneratedKeys && !Boolean.FALSE.equals(generatedKeysRequest);
int objectId = readGeneratedKeys ? session.getNextId() : 0;
synchronized (session) {
int updateCount = 0;
ResultRemote generatedKeys = null;
boolean autoCommit = false;
for (int i = 0, count = 0; i < transferList.size(); i++) {
prepareIfRequired();
......@@ -206,20 +211,21 @@ public class CommandRemote implements CommandInterface {
session.traceOperation("COMMAND_EXECUTE_UPDATE", id);
transfer.writeInt(SessionRemote.COMMAND_EXECUTE_UPDATE).writeInt(id);
sendParameters(transfer);
if (session.getClientVersion() >= Constants.TCP_PROTOCOL_VERSION_17) {
if (Boolean.FALSE.equals(generatedKeys)) {
if (supportsGeneratedKeys) {
if (Boolean.FALSE.equals(generatedKeysRequest)) {
transfer.writeInt(0);
} else if (Boolean.TRUE.equals(generatedKeys)) {
readGeneratedKeys = false;
} else if (Boolean.TRUE.equals(generatedKeysRequest)) {
transfer.writeInt(1);
} else if (generatedKeys instanceof int[]) {
int[] keys = (int[]) generatedKeys;
} else if (generatedKeysRequest instanceof int[]) {
int[] keys = (int[]) generatedKeysRequest;
transfer.writeInt(2);
transfer.writeInt(keys.length);
for (int key : keys) {
transfer.writeInt(key);
}
} else if (generatedKeys instanceof String[]) {
String[] keys = (String[]) generatedKeys;
} else if (generatedKeysRequest instanceof String[]) {
String[] keys = (String[]) generatedKeysRequest;
transfer.writeInt(3);
transfer.writeInt(keys.length);
for (String key : keys) {
......@@ -227,11 +233,20 @@ public class CommandRemote implements CommandInterface {
}
} else {
transfer.writeInt(0);
readGeneratedKeys = false;
}
}
session.done(transfer);
updateCount = transfer.readInt();
autoCommit = transfer.readBoolean();
if (readGeneratedKeys) {
int columnCount = transfer.readInt();
if (generatedKeys != null) {
generatedKeys.close();
generatedKeys = null;
}
generatedKeys = new ResultRemote(session, transfer, objectId, columnCount, Integer.MAX_VALUE);
}
} catch (IOException e) {
session.removeServer(e, i--, ++count);
}
......@@ -239,7 +254,10 @@ public class CommandRemote implements CommandInterface {
session.setAutoCommitFromServer(autoCommit);
session.autoCommitIfCluster();
session.readSessionState();
return updateCount;
if (generatedKeys != null) {
return new ResultWithGeneratedKeys.WithKeys(updateCount, generatedKeys);
}
return ResultWithGeneratedKeys.of(updateCount);
}
}
......
......@@ -1254,8 +1254,8 @@ public class Session extends SessionWithState {
*/
public void setCurrentCommand(Command command, Object generatedKeysRequest) {
this.currentCommand = command;
// Preserve generated keys in case of a new query so they can be read with
// CALL GET_GENERATED_KEYS()
// Preserve generated keys in case of a new query due to possible nested
// queries in update
if (command != null && !command.isQuery()) {
getGeneratedKeys().clear(generatedKeysRequest);
}
......@@ -1809,4 +1809,10 @@ public class Session extends SessionWithState {
public void setColumnNamerConfiguration(ColumnNamerConfiguration columnNamerConfiguration) {
this.columnNamerConfiguration = columnNamerConfiguration;
}
@Override
public boolean isSupportsGeneratedKeys() {
return true;
}
}
......@@ -153,4 +153,13 @@ public interface SessionInterface extends Closeable {
* @return the current schema name
*/
String getCurrentSchemaName();
/**
* Returns is this session supports generated keys.
*
* @return {@code true} if generated keys are supported, {@code false} if only
* {@code SCOPE_IDENTITY()} is supported
*/
boolean isSupportsGeneratedKeys();
}
......@@ -869,4 +869,10 @@ public class SessionRemote extends SessionWithState implements DataHandler {
public void setCurrentSchemaName(String schema) {
throw DbException.getUnsupportedException("setSchema && remote session");
}
@Override
public boolean isSupportsGeneratedKeys() {
return getClientVersion() >= Constants.TCP_PROTOCOL_VERSION_17;
}
}
......@@ -115,7 +115,7 @@ public class Function extends Expression implements FunctionCall {
public static final int DATABASE = 150, USER = 151, CURRENT_USER = 152,
IDENTITY = 153, SCOPE_IDENTITY = 154, AUTOCOMMIT = 155,
READONLY = 156, DATABASE_PATH = 157, LOCK_TIMEOUT = 158,
DISK_SPACE_USED = 159, SIGNAL = 160, GET_GENERATED_KEYS = 161;
DISK_SPACE_USED = 159, SIGNAL = 160;
private static final Pattern SIGNAL_PATTERN = Pattern.compile("[0-9A-Z]{5}");
......@@ -491,7 +491,6 @@ public class Function extends Expression implements FunctionCall {
addFunctionNotDeterministic("DISK_SPACE_USED", DISK_SPACE_USED,
1, Value.LONG);
addFunctionWithNull("SIGNAL", SIGNAL, 2, Value.NULL);
addFunction("GET_GENERATED_KEYS", GET_GENERATED_KEYS, 0, Value.RESULT_SET);
addFunction("H2VERSION", H2VERSION, 0, Value.STRING);
// TableFunction
......@@ -925,9 +924,6 @@ public class Function extends Expression implements FunctionCall {
case DISK_SPACE_USED:
result = ValueLong.get(getDiskSpaceUsed(session, v0));
break;
case GET_GENERATED_KEYS:
result = ValueResultSet.get(session.getGeneratedKeys().getKeys());
break;
case CAST:
case CONVERT: {
v0 = v0.convertTo(dataType);
......
......@@ -85,7 +85,6 @@ public class JdbcConnection extends TraceObject
private CommandInterface getReadOnly, getGeneratedKeys;
private CommandInterface setLockMode, getLockMode;
private CommandInterface setQueryTimeout, getQueryTimeout;
private boolean oldGetGeneratedKeys;
private int savepointId;
private String catalog;
......@@ -1583,21 +1582,6 @@ public class JdbcConnection extends TraceObject
* INTERNAL
*/
ResultSet getGeneratedKeys(JdbcStatement stat, int id) {
if (!oldGetGeneratedKeys) {
try {
getGeneratedKeys = prepareCommand("CALL GET_GENERATED_KEYS()", getGeneratedKeys);
ResultInterface result = getGeneratedKeys.executeQuery(Integer.MAX_VALUE, true);
ResultSet rs = new JdbcResultSet(this, stat, getGeneratedKeys, result,
id, false, true, false);
return rs;
} catch (DbException ex) {
if (ex.getErrorCode() == ErrorCode.FUNCTION_NOT_FOUND_1) {
oldGetGeneratedKeys = true;
} else {
throw ex;
}
}
}
getGeneratedKeys = prepareCommand(
"SELECT SCOPE_IDENTITY() "
+ "WHERE SCOPE_IDENTITY() IS NOT NULL",
......
......@@ -31,6 +31,7 @@ import org.h2.expression.ParameterInterface;
import org.h2.message.DbException;
import org.h2.message.TraceObject;
import org.h2.result.ResultInterface;
import org.h2.result.ResultWithGeneratedKeys;
import org.h2.util.DateTimeUtils;
import org.h2.util.IOUtils;
import org.h2.util.MergedResultSet;
......@@ -195,7 +196,14 @@ public class JdbcPreparedStatement extends JdbcStatement implements
synchronized (session) {
try {
setExecutingStatement(command);
updateCount = command.executeUpdate(generatedKeysRequest);
ResultWithGeneratedKeys result = command.executeUpdate(generatedKeysRequest);
updateCount = result.getUpdateCount();
ResultInterface gk = result.getGeneratedKeys();
if (gk != null) {
int id = getNextId(TraceObject.RESULT_SET);
generatedKeys = new JdbcResultSet(conn, this, command, gk, id,
false, true, false);
}
} finally {
setExecutingStatement(null);
}
......@@ -238,7 +246,13 @@ public class JdbcPreparedStatement extends JdbcStatement implements
updatable, cachedColumnLabelMap);
} else {
returnsResultSet = false;
updateCount = command.executeUpdate(generatedKeysRequest);
ResultWithGeneratedKeys result = command.executeUpdate(generatedKeysRequest);
updateCount = result.getUpdateCount();
ResultInterface gk = result.getGeneratedKeys();
if (gk != null) {
generatedKeys = new JdbcResultSet(conn, this, command, gk, id,
false, true, false);
}
}
} finally {
if (!lazy) {
......@@ -1263,7 +1277,8 @@ public class JdbcPreparedStatement extends JdbcStatement implements
}
try {
result[i] = executeUpdateInternal();
ResultSet rs = conn.getGeneratedKeys(this, id);
// Cannot use own implementation, it returns batch identities
ResultSet rs = super.getGeneratedKeys();
batchIdentities.add(rs);
} catch (Exception re) {
SQLException e = logAndConvert(re);
......
......@@ -18,6 +18,8 @@ import org.h2.engine.SysProperties;
import org.h2.message.DbException;
import org.h2.message.TraceObject;
import org.h2.result.ResultInterface;
import org.h2.result.ResultWithGeneratedKeys;
import org.h2.tools.SimpleResultSet;
import org.h2.util.New;
import org.h2.util.ParserUtil;
import org.h2.util.StringUtils;
......@@ -33,6 +35,7 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
protected int maxRows;
protected int fetchSize = SysProperties.SERVER_RESULT_SET_FETCH_SIZE;
protected int updateCount;
protected JdbcResultSet generatedKeys;
protected final int resultSetType;
protected final int resultSetConcurrency;
protected final boolean closedByResultSet;
......@@ -163,7 +166,14 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
synchronized (session) {
setExecutingStatement(command);
try {
updateCount = command.executeUpdate(generatedKeysRequest);
ResultWithGeneratedKeys result = command.executeUpdate(generatedKeysRequest);
updateCount = result.getUpdateCount();
ResultInterface gk = result.getGeneratedKeys();
if (gk != null) {
int id = getNextId(TraceObject.RESULT_SET);
generatedKeys = new JdbcResultSet(conn, this, command, gk, id,
false, true, false);
}
} finally {
setExecutingStatement(null);
}
......@@ -219,7 +229,13 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
closedByResultSet, scrollable, updatable);
} else {
returnsResultSet = false;
updateCount = command.executeUpdate(generatedKeysRequest);
ResultWithGeneratedKeys result = command.executeUpdate(generatedKeysRequest);
updateCount = result.getUpdateCount();
ResultInterface gk = result.getGeneratedKeys();
if (gk != null) {
generatedKeys = new JdbcResultSet(conn, this, command, gk, id,
false, true, false);
}
}
} finally {
if (!lazy) {
......@@ -816,6 +832,13 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
debugCodeAssign("ResultSet", TraceObject.RESULT_SET, id, "getGeneratedKeys()");
}
checkClosed();
if (generatedKeys != null) {
return generatedKeys;
}
if (session.isSupportsGeneratedKeys()) {
return new SimpleResultSet();
}
// Old server, so use SCOPE_IDENTITY()
return conn.getGeneratedKeys(this, id);
} catch (Exception e) {
throw logAndConvert(e);
......@@ -1197,11 +1220,15 @@ public class JdbcStatement extends TraceObject implements Statement, JdbcStateme
if (resultSet != null) {
resultSet.closeInternal();
}
if (generatedKeys != null) {
generatedKeys.closeInternal();
}
}
} finally {
cancelled = false;
resultSet = null;
updateCount = -1;
generatedKeys = null;
}
}
......
/*
* Copyright 2004-2018 H2 Group. Multiple-Licensed under the MPL 2.0,
* and the EPL 1.0 (http://h2database.com/html/license.html).
* Initial Developer: H2 Group
*/
package org.h2.result;
/**
* Result of update command with optional generated keys.
*/
public class ResultWithGeneratedKeys {
/**
* Result of update command with generated keys;
*/
public static final class WithKeys extends ResultWithGeneratedKeys {
private final ResultInterface generatedKeys;
/**
* Creates a result with update count and generated keys.
*
* @param updateCount
* update count
* @param generatedKeys
* generated keys
*/
public WithKeys(int updateCount, ResultInterface generatedKeys) {
super(updateCount);
this.generatedKeys = generatedKeys;
}
@Override
public ResultInterface getGeneratedKeys() {
return generatedKeys;
}
}
/**
* Returns a result with only update count.
*
* @param updateCount
* update count
*/
public static ResultWithGeneratedKeys of(int updateCount) {
return new ResultWithGeneratedKeys(updateCount);
}
private final int updateCount;
ResultWithGeneratedKeys(int updateCount) {
this.updateCount = updateCount;
}
/**
* Returns generated keys, or {@code null}.
*
* @return generated keys, or {@code null}
*/
public ResultInterface getGeneratedKeys() {
return null;
}
/**
* Returns update count.
*
* @return update count
*/
public int getUpdateCount() {
return updateCount;
}
}
......@@ -31,6 +31,7 @@ import org.h2.jdbc.JdbcSQLException;
import org.h2.message.DbException;
import org.h2.result.ResultColumn;
import org.h2.result.ResultInterface;
import org.h2.result.ResultWithGeneratedKeys;
import org.h2.store.LobStorageInterface;
import org.h2.util.IOUtils;
import org.h2.util.SmallLRUCache;
......@@ -353,12 +354,15 @@ public class TcpServerThread implements Runnable {
int id = transfer.readInt();
Command command = (Command) cache.getObject(id, false);
setParameters(command);
boolean supportsGeneratedKeys = clientVersion >= Constants.TCP_PROTOCOL_VERSION_17;
boolean writeGeneratedKeys = supportsGeneratedKeys;
Object generatedKeysRequest;
if (clientVersion >= Constants.TCP_PROTOCOL_VERSION_17) {
if (supportsGeneratedKeys) {
int type = transfer.readInt();
switch (type) {
default:
generatedKeysRequest = false;
writeGeneratedKeys = false;
break;
case 1:
generatedKeysRequest = true;
......@@ -382,12 +386,12 @@ public class TcpServerThread implements Runnable {
}
}
} else {
generatedKeysRequest = true;
generatedKeysRequest = false;
}
int old = session.getModificationId();
int updateCount;
ResultWithGeneratedKeys result;
synchronized (session) {
updateCount = command.executeUpdate(generatedKeysRequest);
result = command.executeUpdate(generatedKeysRequest);
}
int status;
if (session.isClosed()) {
......@@ -396,8 +400,22 @@ public class TcpServerThread implements Runnable {
} else {
status = getState(old);
}
transfer.writeInt(status).writeInt(updateCount).
transfer.writeInt(status).writeInt(result.getUpdateCount()).
writeBoolean(session.getAutoCommit());
if (writeGeneratedKeys) {
ResultInterface generatedKeys = result.getGeneratedKeys();
int columnCount = generatedKeys.getVisibleColumnCount();
transfer.writeInt(columnCount);
int rowCount = generatedKeys.getRowCount();
transfer.writeInt(rowCount);
for (int i = 0; i < columnCount; i++) {
ResultColumn.writeColumn(transfer, generatedKeys, i);
}
for (int i = 0; i < rowCount; i++) {
sendRow(generatedKeys);
}
generatedKeys.close();
}
transfer.flush();
break;
}
......
......@@ -1355,7 +1355,7 @@ public class TestPreparedStatement extends TestBase {
"insert into test values (30), (next value for seq),"
+ " (next value for seq), (next value for seq), (20)",
PreparedStatement.RETURN_GENERATED_KEYS);
prep.executeUpdate();
prep.executeUpdate();
ResultSet rs = prep.getGeneratedKeys();
rs.next();
assertEquals(1L, rs.getLong(1));
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论