Index: src/com/ibatis/sqlmap/engine/execution/SqlExecutor.java =================================================================== --- src/com/ibatis/sqlmap/engine/execution/SqlExecutor.java (revision 756020) +++ src/com/ibatis/sqlmap/engine/execution/SqlExecutor.java (working copy) @@ -19,8 +19,10 @@ import com.ibatis.sqlmap.engine.mapping.parameter.ParameterMap; import com.ibatis.sqlmap.engine.mapping.result.ResultMap; import com.ibatis.sqlmap.engine.mapping.result.ResultObjectFactoryUtil; +import com.ibatis.sqlmap.engine.mapping.statement.InsertStatement; import com.ibatis.sqlmap.engine.mapping.statement.MappedStatement; import com.ibatis.sqlmap.engine.mapping.statement.RowHandlerCallback; +import com.ibatis.sqlmap.engine.mapping.statement.SelectKeyStatement; import com.ibatis.sqlmap.engine.scope.ErrorContext; import com.ibatis.sqlmap.engine.scope.StatementScope; import com.ibatis.sqlmap.engine.scope.SessionScope; @@ -71,13 +73,51 @@ setupResultObjectFactory(statementScope); int rows = 0; try { + MappedStatement mappedStatement = statementScope.getStatement(); + String[] keyColumns = new String[0]; + if (mappedStatement instanceof InsertStatement) + { + InsertStatement insertStmt = (InsertStatement) mappedStatement; + String generatedKeyColumn = insertStmt.getSelectKeyStatement().getGeneratedKeyColumn(); + if (generatedKeyColumn != null) + { + keyColumns = new String[] {generatedKeyColumn}; + } + } + errorContext.setMoreInfo("Check the SQL Statement (preparation failed)."); - ps = prepareStatement(statementScope.getSession(), conn, sql); + ps = prepareStatement(statementScope.getSession(), conn, sql, keyColumns); setStatementTimeout(statementScope.getStatement(), ps); errorContext.setMoreInfo("Check the parameters (set parameters failed)."); statementScope.getParameterMap().setParameters(statementScope, ps, parameters); errorContext.setMoreInfo("Check the statement (update failed)."); + // Reset any generate key + SelectKeyStatement.GeneratedKeyThreadLocal.clear(); ps.execute(); + // Retrieve generated key, if necessary + if (keyColumns.length > 0) + { + ResultSet rsKeys = ps.getGeneratedKeys(); + try + { + int rsColumnCount = rsKeys.getMetaData().getColumnCount(); + if (rsColumnCount != 1) + { + throw new SQLException("Multi-column generated keys are not supported."); + } + if (!rsKeys.next()) + { + throw new SQLException("No results returned for generated keys."); + } + Object generatedKey = rsKeys.getObject(keyColumns[0]); + // Store this for later + SelectKeyStatement.GeneratedKeyThreadLocal.setKey(generatedKey); + } + finally + { + try { rsKeys.close(); } catch (Throwable e) {} + } + } rows = ps.getUpdateCount(); } finally { closeStatement(statementScope.getSession(), ps); @@ -500,6 +540,21 @@ } } + private static PreparedStatement prepareStatement(SessionScope sessionScope, Connection conn, String sql, String[] returnColumns) throws SQLException { + SqlMapExecutorDelegate delegate = ((SqlMapClientImpl) sessionScope.getSqlMapExecutor()).getDelegate(); + if (sessionScope.hasPreparedStatementFor(sql)) { + return sessionScope.getPreparedStatement((sql)); + } else { + PreparedStatement ps; + if (returnColumns == null || returnColumns.length == 0) + ps = conn.prepareStatement(sql); + else + ps = conn.prepareStatement(sql, returnColumns); + sessionScope.putPreparedStatement(delegate, sql, ps); + return ps; + } + } + private CallableStatement prepareCall(SessionScope sessionScope, Connection conn, String sql) throws SQLException { SqlMapExecutorDelegate delegate = ((SqlMapClientImpl) sessionScope.getSqlMapExecutor()).getDelegate(); if (sessionScope.hasPreparedStatementFor(sql)) { Index: src/com/ibatis/sqlmap/engine/config/MappedStatementConfig.java =================================================================== --- src/com/ibatis/sqlmap/engine/config/MappedStatementConfig.java (revision 756020) +++ src/com/ibatis/sqlmap/engine/config/MappedStatementConfig.java (working copy) @@ -114,7 +114,7 @@ delegate.addMappedStatement(mappedStatement); } - public void setSelectKeyStatement(SqlSource processor, String resultClassName, String keyPropName, boolean runAfterSQL, String type) { + public void setSelectKeyStatement(SqlSource processor, String resultClassName, String keyPropName, boolean runAfterSQL, String type, String generatedKeyColumn) { if (rootStatement instanceof InsertStatement) { InsertStatement insertStatement = ((InsertStatement) rootStatement); Class parameterClass = insertStatement.getParameterClass(); @@ -129,6 +129,7 @@ selectKeyStatement.setResource(errorContext.getResource()); selectKeyStatement.setKeyProperty(keyPropName); selectKeyStatement.setRunAfterSQL(runAfterSQL); + selectKeyStatement.setGeneratedKeyColumn(generatedKeyColumn); // process the type (pre or post) attribute if (type != null) { selectKeyStatement.setRunAfterSQL("post".equals(type)); Index: src/com/ibatis/sqlmap/engine/builder/xml/SqlStatementParser.java =================================================================== --- src/com/ibatis/sqlmap/engine/builder/xml/SqlStatementParser.java (revision 756020) +++ src/com/ibatis/sqlmap/engine/builder/xml/SqlStatementParser.java (working copy) @@ -5,9 +5,9 @@ import com.ibatis.sqlmap.engine.config.*; import com.ibatis.sqlmap.engine.mapping.statement.*; import com.ibatis.sqlmap.client.*; -import org.w3c.dom.CharacterData; import org.w3c.dom.*; +import java.sql.SQLException; import java.util.Properties; public class SqlStatementParser { @@ -108,7 +108,19 @@ String keyPropName = attributes.getProperty("keyProperty"); String resultClassName = attributes.getProperty("resultClass"); String type = attributes.getProperty("type"); - config.setSelectKeyStatement(new XMLSqlSource(state, child), resultClassName, keyPropName, foundSQLFirst, type); + String selectQuery = ((CharacterData) child.getFirstChild( )).getData(); + String generatedKeyColumn = null; + if (selectQuery.trim().toUpperCase().startsWith("KEY_COLUMN:")) + { + String sql = selectQuery.trim(); + generatedKeyColumn = sql.substring(11); + if (generatedKeyColumn.length() == 0) + { + throw new SqlMapException("Invalid key column specification: " + selectQuery.trim()); + } + } + boolean useGeneratedKey = selectQuery.trim().toUpperCase().startsWith("KEY_COLUMN:"); + config.setSelectKeyStatement(new XMLSqlSource(state, child), resultClassName, keyPropName, foundSQLFirst, type, generatedKeyColumn); break; } } Index: src/com/ibatis/sqlmap/engine/mapping/statement/SelectKeyStatement.java =================================================================== --- src/com/ibatis/sqlmap/engine/mapping/statement/SelectKeyStatement.java (revision 756020) +++ src/com/ibatis/sqlmap/engine/mapping/statement/SelectKeyStatement.java (working copy) @@ -15,10 +15,20 @@ */ package com.ibatis.sqlmap.engine.mapping.statement; +import com.ibatis.common.beans.Probe; +import com.ibatis.common.beans.ProbeFactory; +import com.ibatis.common.resources.Resources; import com.ibatis.sqlmap.client.event.RowHandler; +import com.ibatis.sqlmap.engine.exchange.DataExchange; +import com.ibatis.sqlmap.engine.mapping.result.ResultMap; +import com.ibatis.sqlmap.engine.mapping.result.ResultMapping; import com.ibatis.sqlmap.engine.scope.StatementScope; import com.ibatis.sqlmap.engine.transaction.Transaction; +import com.ibatis.sqlmap.engine.type.TypeHandler; +import com.ibatis.sqlmap.engine.type.TypeHandlerFactory; +import java.beans.BeanDescriptor; +import java.sql.Connection; import java.sql.SQLException; import java.util.List; @@ -26,6 +36,7 @@ private String keyProperty; private boolean runAfterSQL; + private String generatedKeyColumn; public String getKeyProperty() { return keyProperty; @@ -43,6 +54,21 @@ this.runAfterSQL = runAfterSQL; } + public boolean isUseGeneratedKey() + { + return generatedKeyColumn != null; + } + + public String getGeneratedKeyColumn() + { + return this.generatedKeyColumn; + } + + public void setGeneratedKeyColumn(String generatedKeyColumn) + { + this.generatedKeyColumn = generatedKeyColumn; + } + public List executeQueryForList(StatementScope statementScope, Transaction trans, Object parameterObject, int skipResults, int maxResults) throws SQLException { throw new SQLException("Select Key statements cannot be executed for a list."); @@ -53,4 +79,59 @@ throw new SQLException("Select Key statements cannot be executed with a row handler."); } + @SuppressWarnings("unchecked") + protected void executeQueryWithCallback( + StatementScope statementScope, Connection conn, Object parameterObject, Object resultObject, + RowHandler rowHandler, int skipResults, + int maxResults) throws SQLException + { + if (generatedKeyColumn == null) + { + super.executeQueryWithCallback(statementScope, conn, parameterObject, resultObject, rowHandler, skipResults, maxResults); + } + else + { + // We are using generate keys, so spoof a resultset + Object generatedKey = GeneratedKeyThreadLocal.getKey(); + if (generatedKey == null) + { + throw new SQLException("No generated key value available for column: " + this.generatedKeyColumn); + } + // Cast the generated key to the correct type on the resultObject + ResultMap resultMap = getSql().getResultMap(statementScope, parameterObject); + Class resultType = resultMap.getResultClass(); + + if (resultType.isAssignableFrom(generatedKey.getClass())) + { + // No conversion required + } + else + { + // This is untidy, but the type conversion APIs were not designed for this + TypeHandler resultTypeHandler = new TypeHandlerFactory().getTypeHandler(resultType); + generatedKey = resultTypeHandler.valueOf(generatedKey.toString()); + } + rowHandler.handleRow(generatedKey); + } + } + + public static class GeneratedKeyThreadLocal + { + private static final ThreadLocal tl = new ThreadLocal(); + + public static final Object getKey( ) + { + return tl.get(); + } + + public static final void setKey(Object key) + { + tl.set(key); + } + + public static final void clear( ) + { + tl.remove(); + } + } }