Repository: metamodel
Updated Branches:
  refs/heads/5.x c51994c1f -> 95226631f


METAMODEL-6: Added UpdateSummary building for JDBC module

Project: http://git-wip-us.apache.org/repos/asf/metamodel/repo
Commit: http://git-wip-us.apache.org/repos/asf/metamodel/commit/d888907b
Tree: http://git-wip-us.apache.org/repos/asf/metamodel/tree/d888907b
Diff: http://git-wip-us.apache.org/repos/asf/metamodel/diff/d888907b

Branch: refs/heads/5.x
Commit: d888907b3b27d821e0b8fd9291987e9b56e44c68
Parents: c51994c
Author: Kasper Sørensen <[email protected]>
Authored: Sat Jan 2 21:59:14 2016 +0100
Committer: Kasper Sørensen <[email protected]>
Committed: Sat Jan 2 21:59:14 2016 +0100

----------------------------------------------------------------------
 .../metamodel/AbstractUpdateCallback.java       |  12 +-
 .../apache/metamodel/UpdateSummaryBuilder.java  | 130 +++++++++++++++++++
 .../metamodel/jdbc/JdbcBatchUpdateCallback.java |   8 +-
 .../metamodel/jdbc/JdbcDeleteBuilder.java       |   4 +-
 .../metamodel/jdbc/JdbcDropTableBuilder.java    |   4 +-
 .../metamodel/jdbc/JdbcInsertBuilder.java       |   4 +-
 .../jdbc/JdbcSimpleUpdateCallback.java          |  11 +-
 .../metamodel/jdbc/JdbcUpdateBuilder.java       |   4 +-
 .../metamodel/jdbc/JdbcUpdateCallback.java      | 108 +++++++++++----
 .../metamodel/jdbc/JdbcTestTemplates.java       |   4 +-
 .../jdbc/integrationtests/PostgresqlTest.java   |  39 ++++++
 11 files changed, 287 insertions(+), 41 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/core/src/main/java/org/apache/metamodel/AbstractUpdateCallback.java
----------------------------------------------------------------------
diff --git 
a/core/src/main/java/org/apache/metamodel/AbstractUpdateCallback.java 
b/core/src/main/java/org/apache/metamodel/AbstractUpdateCallback.java
index 20842a6..a1c2f1a 100644
--- a/core/src/main/java/org/apache/metamodel/AbstractUpdateCallback.java
+++ b/core/src/main/java/org/apache/metamodel/AbstractUpdateCallback.java
@@ -48,14 +48,14 @@ public abstract class AbstractUpdateCallback implements 
UpdateCallback {
     }
 
     @Override
-    public TableDropBuilder dropTable(String schemaName, String tableName) 
throws IllegalArgumentException,
+    public final TableDropBuilder dropTable(String schemaName, String 
tableName) throws IllegalArgumentException,
             IllegalStateException, UnsupportedOperationException {
         final Table table = getTable(schemaName, tableName);
         return dropTable(table);
     }
 
     @Override
-    public TableDropBuilder dropTable(Schema schema, String tableName) throws 
IllegalArgumentException,
+    public final TableDropBuilder dropTable(Schema schema, String tableName) 
throws IllegalArgumentException,
             IllegalStateException, UnsupportedOperationException {
         final Table table = schema.getTableByName(tableName);
         if (table == null) {
@@ -72,7 +72,7 @@ public abstract class AbstractUpdateCallback implements 
UpdateCallback {
     }
 
     @Override
-    public RowInsertionBuilder insertInto(String schemaName, String tableName) 
throws IllegalArgumentException,
+    public final RowInsertionBuilder insertInto(String schemaName, String 
tableName) throws IllegalArgumentException,
             IllegalStateException, UnsupportedOperationException {
         return insertInto(getTable(schemaName, tableName));
     }
@@ -101,7 +101,7 @@ public abstract class AbstractUpdateCallback implements 
UpdateCallback {
     }
 
     @Override
-    public RowDeletionBuilder deleteFrom(String schemaName, String tableName) 
throws IllegalArgumentException,
+    public final RowDeletionBuilder deleteFrom(String schemaName, String 
tableName) throws IllegalArgumentException,
             IllegalStateException, UnsupportedOperationException {
         final Table table = getTable(schemaName, tableName);
         return deleteFrom(table);
@@ -138,7 +138,7 @@ public abstract class AbstractUpdateCallback implements 
UpdateCallback {
 
     @Override
     public boolean isInsertSupported() {
-        // since 2.0 all updateable datacontext have create table support
+        // since 2.0 all updateable datacontext have insert into table support
         return true;
     }
 
@@ -148,7 +148,7 @@ public abstract class AbstractUpdateCallback implements 
UpdateCallback {
     }
 
     @Override
-    public RowUpdationBuilder update(String schemaName, String tableName) 
throws IllegalArgumentException,
+    public final RowUpdationBuilder update(String schemaName, String 
tableName) throws IllegalArgumentException,
             IllegalStateException, UnsupportedOperationException {
         final Table table = getTable(schemaName, tableName);
         return update(table);

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/core/src/main/java/org/apache/metamodel/UpdateSummaryBuilder.java
----------------------------------------------------------------------
diff --git a/core/src/main/java/org/apache/metamodel/UpdateSummaryBuilder.java 
b/core/src/main/java/org/apache/metamodel/UpdateSummaryBuilder.java
new file mode 100644
index 0000000..6469a1b
--- /dev/null
+++ b/core/src/main/java/org/apache/metamodel/UpdateSummaryBuilder.java
@@ -0,0 +1,130 @@
+/**
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.metamodel;
+
+import java.util.HashSet;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+/**
+ * A builder object for {@link UpdateSummary}.
+ */
+public class UpdateSummaryBuilder {
+
+    private int _inserts;
+    private int _updates;
+    private int _deletes;
+    private Set<Object> _generatedKeys;
+
+    public UpdateSummaryBuilder() {
+    }
+
+    public UpdateSummary build() {
+        final Integer insertedRows = (_inserts == -1 ? null : _inserts);
+        final Integer updatedRows = (_updates == -1 ? null : _updates);
+        final Integer deletedRows = (_deletes == -1 ? null : _deletes);
+        final Iterable<Object> generatedKeys;
+        if (_generatedKeys != null) {
+            generatedKeys = new LinkedHashSet<>(_generatedKeys);
+        } else {
+            generatedKeys = null;
+        }
+        return new DefaultUpdateSummary(insertedRows, updatedRows, 
deletedRows, generatedKeys);
+    }
+
+    public UpdateSummaryBuilder addInsert() {
+        return addInserts(1);
+    }
+
+    public UpdateSummaryBuilder addInserts(int inserts) {
+        if (_inserts != -1) {
+            _inserts += inserts;
+        }
+        return this;
+    }
+
+    public UpdateSummaryBuilder makeInsertsUnknown() {
+        _inserts = -1;
+        return this;
+    }
+
+    public UpdateSummaryBuilder addUpdate() {
+        return addUpdates(1);
+    }
+
+    public UpdateSummaryBuilder addUpdates(int updates) {
+        if (_updates != -1) {
+            _updates += updates;
+        }
+        return this;
+    }
+
+    public UpdateSummaryBuilder makeUpdatesUnknown() {
+        _updates = -1;
+        return this;
+    }
+
+    public UpdateSummaryBuilder addDelete() {
+        return addDeletes(1);
+    }
+
+    public UpdateSummaryBuilder addDeletes(int deletes) {
+        if (_deletes != -1) {
+            _deletes += deletes;
+        }
+        return this;
+    }
+
+    public UpdateSummaryBuilder makeDeletesUnknown() {
+        _deletes = -1;
+        return this;
+    }
+
+    public UpdateSummaryBuilder addGeneratedKey(Object key) {
+        if (_generatedKeys == null) {
+            _generatedKeys = new HashSet<>();
+        }
+        _generatedKeys.add(key);
+        return this;
+    }
+
+    public UpdateSummaryBuilder addGeneratedKeys(Object... keys) {
+        if (_generatedKeys == null) {
+            _generatedKeys = new HashSet<>();
+        }
+        for (Object key : keys) {
+            _generatedKeys.add(key);
+        }
+        return this;
+    }
+
+    public UpdateSummaryBuilder addGeneratedKeys(Iterable<?> keys) {
+        if (_generatedKeys == null) {
+            _generatedKeys = new HashSet<>();
+        }
+        for (Object key : keys) {
+            _generatedKeys.add(key);
+        }
+        return this;
+    }
+
+    public UpdateSummaryBuilder makeGeneratedKeysUnknown() {
+        _generatedKeys = null;
+        return this;
+    }
+}

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcBatchUpdateCallback.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcBatchUpdateCallback.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcBatchUpdateCallback.java
index 100c5d0..52478fc 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcBatchUpdateCallback.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcBatchUpdateCallback.java
@@ -68,7 +68,13 @@ final class JdbcBatchUpdateCallback extends 
JdbcUpdateCallback {
     }
 
     @Override
-    protected void executePreparedStatement(PreparedStatement st) throws 
SQLException {
+    protected int executePreparedStatement(PreparedStatement st) throws 
SQLException {
         st.addBatch();
+        return -1;
+    }
+    
+    @Override
+    protected boolean isGeneratedKeysCollectionEnabled() {
+        return false;
     }
 }

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDeleteBuilder.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDeleteBuilder.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDeleteBuilder.java
index bf6aaf2..98d0245 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDeleteBuilder.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDeleteBuilder.java
@@ -62,7 +62,7 @@ final class JdbcDeleteBuilder extends 
AbstractRowDeletionBuilder {
 
         logger.debug("Delete statement created: {}", sql);
         final boolean reuseStatement = !_inlineValues;
-        final PreparedStatement st = _updateCallback.getPreparedStatement(sql, 
reuseStatement);
+        final PreparedStatement st = _updateCallback.getPreparedStatement(sql, 
reuseStatement, false);
         try {
             if (reuseStatement) {
                 int valueCounter = 1;
@@ -75,7 +75,7 @@ final class JdbcDeleteBuilder extends 
AbstractRowDeletionBuilder {
                     }
                 }
             }
-            _updateCallback.executePreparedStatement(st, reuseStatement);
+            _updateCallback.executeDelete(st, reuseStatement);
         } catch (SQLException e) {
             throw JdbcUtils.wrapException(e, "execute delete statement: " + 
sql);
         } finally {

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDropTableBuilder.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDropTableBuilder.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDropTableBuilder.java
index e4f387b..87ae31d 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDropTableBuilder.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcDropTableBuilder.java
@@ -46,9 +46,9 @@ final class JdbcDropTableBuilder extends 
AbstractTableDropBuilder implements Tab
     @Override
     public void execute() {
         final String sql = createSqlStatement();
-        final PreparedStatement statement = 
_updateCallback.getPreparedStatement(sql, false);
+        final PreparedStatement statement = 
_updateCallback.getPreparedStatement(sql, false, false);
         try {
-            _updateCallback.executePreparedStatement(statement, false);
+            _updateCallback.executePreparedStatement(statement, false, false);
 
             // remove the table reference from the schema
             final Schema schema = getTable().getSchema();

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcInsertBuilder.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcInsertBuilder.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcInsertBuilder.java
index 66cdbe7..5b17c9a 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcInsertBuilder.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcInsertBuilder.java
@@ -66,7 +66,7 @@ final class JdbcInsertBuilder extends 
AbstractRowInsertionBuilder<JdbcUpdateCall
                }
                final JdbcUpdateCallback updateCallback = getUpdateCallback();
                final boolean reuseStatement = !_inlineValues;
-               final PreparedStatement st = 
updateCallback.getPreparedStatement(sql, reuseStatement);
+               final PreparedStatement st = 
updateCallback.getPreparedStatement(sql, reuseStatement, true);
                try {
                        if (reuseStatement) {
                                Column[] columns = getColumns();
@@ -81,7 +81,7 @@ final class JdbcInsertBuilder extends 
AbstractRowInsertionBuilder<JdbcUpdateCall
                                        }
                                }
                        }
-                       updateCallback.executePreparedStatement(st, 
reuseStatement);
+                       updateCallback.executeInsert(st, reuseStatement);
                } catch (SQLException e) {
                        throw JdbcUtils.wrapException(e, "execute insert 
statement: " + sql);
                } finally {

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcSimpleUpdateCallback.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcSimpleUpdateCallback.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcSimpleUpdateCallback.java
index 3d44e42..6d174bf 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcSimpleUpdateCallback.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcSimpleUpdateCallback.java
@@ -34,14 +34,19 @@ final class JdbcSimpleUpdateCallback extends 
JdbcUpdateCallback {
     public JdbcSimpleUpdateCallback(JdbcDataContext dataContext) {
         super(dataContext);
     }
-    
+
     @Override
     protected void closePreparedStatement(PreparedStatement preparedStatement) 
{
         FileHelper.safeClose(preparedStatement);
     }
 
     @Override
-    protected void executePreparedStatement(PreparedStatement st) throws 
SQLException {
-        st.executeUpdate();
+    protected int executePreparedStatement(PreparedStatement st) throws 
SQLException {
+        return st.executeUpdate();
+    }
+    
+    @Override
+    protected boolean isGeneratedKeysCollectionEnabled() {
+        return true;
     }
 }

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateBuilder.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateBuilder.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateBuilder.java
index 6610a7e..bcdb810 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateBuilder.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateBuilder.java
@@ -62,7 +62,7 @@ final class JdbcUpdateBuilder extends 
AbstractRowUpdationBuilder {
         String sql = createSqlStatement();
         logger.debug("Update statement created: {}", sql);
         final boolean reuseStatement = !_inlineValues;
-        final PreparedStatement st = _updateCallback.getPreparedStatement(sql, 
reuseStatement);
+        final PreparedStatement st = _updateCallback.getPreparedStatement(sql, 
reuseStatement, false);
         try {
             if (reuseStatement) {
                 Column[] columns = getColumns();
@@ -90,7 +90,7 @@ final class JdbcUpdateBuilder extends 
AbstractRowUpdationBuilder {
                     }
                 }
             }
-            _updateCallback.executePreparedStatement(st, reuseStatement);
+            _updateCallback.executeUpdate(st, reuseStatement);
         } catch (SQLException e) {
             throw JdbcUtils.wrapException(e, "execute update statement: " + 
sql);
         } finally {

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateCallback.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateCallback.java 
b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateCallback.java
index 5e60826..29fc263 100644
--- a/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateCallback.java
+++ b/jdbc/src/main/java/org/apache/metamodel/jdbc/JdbcUpdateCallback.java
@@ -20,10 +20,15 @@ package org.apache.metamodel.jdbc;
 
 import java.sql.Connection;
 import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.sql.Statement;
 
 import org.apache.metamodel.AbstractUpdateCallback;
 import org.apache.metamodel.UpdateCallback;
+import org.apache.metamodel.UpdateSummary;
+import org.apache.metamodel.UpdateSummaryBuilder;
 import org.apache.metamodel.create.TableCreationBuilder;
 import org.apache.metamodel.delete.RowDeletionBuilder;
 import org.apache.metamodel.drop.TableDropBuilder;
@@ -31,6 +36,7 @@ import org.apache.metamodel.insert.RowInsertionBuilder;
 import org.apache.metamodel.schema.Schema;
 import org.apache.metamodel.schema.Table;
 import org.apache.metamodel.update.RowUpdationBuilder;
+import org.apache.metamodel.util.FileHelper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -38,32 +44,60 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
 
     private static final Logger logger = 
LoggerFactory.getLogger(JdbcUpdateCallback.class);
 
-    private final JdbcDataContext _dataContext;
     private Connection _connection;
     private String _preparedStatementSql;
     private PreparedStatement _preparedStatement;
+    private final UpdateSummaryBuilder _updateSummaryBuilder;
 
     public JdbcUpdateCallback(JdbcDataContext dataContext) {
         super(dataContext);
-        _dataContext = dataContext;
+        _updateSummaryBuilder = new UpdateSummaryBuilder();
     }
 
+    protected final UpdateSummaryBuilder getUpdateSummaryBuilder() {
+        return _updateSummaryBuilder;
+    }
+    
+    protected abstract boolean isGeneratedKeysCollectionEnabled();
+
     protected abstract void closePreparedStatement(PreparedStatement 
preparedStatement);
 
-    protected abstract void executePreparedStatement(PreparedStatement 
preparedStatement) throws SQLException;
+    protected abstract int executePreparedStatement(PreparedStatement 
preparedStatement) throws SQLException;
+
+    public int executePreparedStatement(PreparedStatement preparedStatement, 
boolean reusedStatement,
+            boolean collectGeneratedKeys) throws SQLException {
+        final int result = executePreparedStatement(preparedStatement);
+
+        if (collectGeneratedKeys && isGeneratedKeysCollectionEnabled()) {
+            try {
+                final ResultSet generatedKeysResultSet = 
preparedStatement.getGeneratedKeys();
+                try {
+                    while (generatedKeysResultSet.next()) {
+                        final Object key = generatedKeysResultSet.getObject(1);
+                        getUpdateSummaryBuilder().addGeneratedKey(key);
+                    }
+                } finally {
+                    FileHelper.safeClose(generatedKeysResultSet);
+                }
+            } catch (SQLFeatureNotSupportedException e) {
+                // ignore
+                logger.debug("Getting generated keys from JDBC statement is 
not supported by driver: {}", e
+                        .getMessage(), e);
+            } catch (SQLException | RuntimeException e) {
+                logger.warn("Failed to get generated keys from JDBC statement: 
{}", e.getMessage(), e);
+            }
+        }
 
-    public void executePreparedStatement(PreparedStatement preparedStatement, 
boolean reusedStatement)
-            throws SQLException {
-        executePreparedStatement(preparedStatement);
         if (!reusedStatement) {
             closePreparedStatement(preparedStatement);
         }
+        return result;
     }
 
     protected final Connection getConnection() {
         if (_connection == null) {
             _connection = getDataContext().getConnection();
-            if (_dataContext.getQueryRewriter().isTransactional()) {
+            if (getDataContext().getQueryRewriter().isTransactional()) {
                 try {
                     _connection.setAutoCommit(false);
                 } catch (SQLException e) {
@@ -80,11 +114,11 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
                 closePreparedStatement(_preparedStatement);
             }
 
-            if (_dataContext.getQueryRewriter().isTransactional()) {
+            if (getDataContext().getQueryRewriter().isTransactional()) {
                 try {
                     commitOrRollback(success);
 
-                    if (_dataContext.isDefaultAutoCommit()) {
+                    if (getDataContext().isDefaultAutoCommit()) {
                         try {
                             getConnection().setAutoCommit(true);
                         } catch (SQLException e) {
@@ -122,19 +156,20 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
 
     @Override
     public final RowInsertionBuilder insertInto(Table table) throws 
IllegalArgumentException, IllegalStateException {
-        return new JdbcInsertBuilder(this, table, 
_dataContext.getQueryRewriter());
+        return new JdbcInsertBuilder(this, table, 
getDataContext().getQueryRewriter());
     }
 
+    // override the return type to the more specific subtype.
     @Override
     public final JdbcDataContext getDataContext() {
-        return _dataContext;
+        return (JdbcDataContext) super.getDataContext();
     }
 
     protected String quoteIfNescesary(String identifier) {
         if (identifier == null) {
             return null;
         }
-        final String quote = _dataContext.getIdentifierQuoteString();
+        final String quote = getDataContext().getIdentifierQuoteString();
         if (quote == null) {
             return identifier;
         }
@@ -153,7 +188,8 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
         return identifier;
     }
 
-    public final PreparedStatement getPreparedStatement(String sql, boolean 
reuseStatement) {
+    public final PreparedStatement getPreparedStatement(String sql, boolean 
reuseStatement,
+            boolean returnGeneratedKeys) {
         final PreparedStatement preparedStatement;
         if (reuseStatement) {
             if (sql.equals(_preparedStatementSql)) {
@@ -167,20 +203,27 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
                         throw e;
                     }
                 }
-                preparedStatement = createPreparedStatement(sql);
+                preparedStatement = createPreparedStatement(sql, 
returnGeneratedKeys);
                 _preparedStatement = preparedStatement;
                 _preparedStatementSql = sql;
             }
         } else {
-            preparedStatement = createPreparedStatement(sql);
+            preparedStatement = createPreparedStatement(sql, 
returnGeneratedKeys);
         }
         return preparedStatement;
     }
-
-    private final PreparedStatement createPreparedStatement(String sql) {
+    
+    private final PreparedStatement createPreparedStatement(String sql, 
boolean returnGeneratedKeys) {
         try {
+            if (returnGeneratedKeys && isGeneratedKeysCollectionEnabled()) {
+                return getConnection().prepareStatement(sql, 
Statement.RETURN_GENERATED_KEYS);
+            }
             return getConnection().prepareStatement(sql);
         } catch (SQLException e) {
+            if (returnGeneratedKeys) {
+                // not all databases support the RETURN_GENERATED_KEYS flag, 
so retry without
+                return createPreparedStatement(sql, false);
+            }
             throw JdbcUtils.wrapException(e, "create prepared statement for: " 
+ sql);
         }
     }
@@ -191,9 +234,9 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
     }
 
     @Override
-    public RowDeletionBuilder deleteFrom(Table table) throws 
IllegalArgumentException, IllegalStateException,
+    public final RowDeletionBuilder deleteFrom(Table table) throws 
IllegalArgumentException, IllegalStateException,
             UnsupportedOperationException {
-        return new JdbcDeleteBuilder(this, table, 
_dataContext.getQueryRewriter());
+        return new JdbcDeleteBuilder(this, table, 
getDataContext().getQueryRewriter());
     }
 
     @Override
@@ -204,7 +247,7 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
     @Override
     public TableDropBuilder dropTable(Table table) throws 
IllegalArgumentException, IllegalStateException,
             UnsupportedOperationException {
-        return new JdbcDropTableBuilder(this, table, 
_dataContext.getQueryRewriter());
+        return new JdbcDropTableBuilder(this, table, 
getDataContext().getQueryRewriter());
     }
 
     @Override
@@ -213,8 +256,29 @@ abstract class JdbcUpdateCallback extends 
AbstractUpdateCallback implements Upda
     }
 
     @Override
-    public RowUpdationBuilder update(Table table) throws 
IllegalArgumentException, IllegalStateException,
+    public final RowUpdationBuilder update(Table table) throws 
IllegalArgumentException, IllegalStateException,
             UnsupportedOperationException {
-        return new JdbcUpdateBuilder(this, table, 
_dataContext.getQueryRewriter());
+        return new JdbcUpdateBuilder(this, table, 
getDataContext().getQueryRewriter());
+    }
+
+    public void executeInsert(PreparedStatement st, boolean reuseStatement) 
throws SQLException {
+        executePreparedStatement(st, reuseStatement, true);
+        _updateSummaryBuilder.addInsert();
+    }
+
+    public void executeUpdate(PreparedStatement st, boolean reuseStatement) 
throws SQLException {
+        final int updates = executePreparedStatement(st, reuseStatement, 
false);
+        _updateSummaryBuilder.addUpdates(updates);
     }
+
+    public void executeDelete(PreparedStatement st, boolean reuseStatement) 
throws SQLException {
+        final int deletes = executePreparedStatement(st, reuseStatement, 
false);
+        _updateSummaryBuilder.addDeletes(deletes);
+    }
+
+    @Override
+    public UpdateSummary getUpdateSummary() {
+        return _updateSummaryBuilder.build();
+    }
+
 }

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/test/java/org/apache/metamodel/jdbc/JdbcTestTemplates.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/test/java/org/apache/metamodel/jdbc/JdbcTestTemplates.java 
b/jdbc/src/test/java/org/apache/metamodel/jdbc/JdbcTestTemplates.java
index f7d0cf2..c4dd018 100644
--- a/jdbc/src/test/java/org/apache/metamodel/jdbc/JdbcTestTemplates.java
+++ b/jdbc/src/test/java/org/apache/metamodel/jdbc/JdbcTestTemplates.java
@@ -40,6 +40,7 @@ import java.util.concurrent.TimeUnit;
 import org.apache.metamodel.BatchUpdateScript;
 import org.apache.metamodel.UpdateCallback;
 import org.apache.metamodel.UpdateScript;
+import org.apache.metamodel.UpdateSummary;
 import org.apache.metamodel.create.ColumnCreationBuilder;
 import org.apache.metamodel.create.CreateTable;
 import org.apache.metamodel.create.TableCreationBuilder;
@@ -73,7 +74,7 @@ public class JdbcTestTemplates {
 
         final Map<Object, Object> map = new HashMap<Object, Object>();
         try {
-            dc.executeUpdate(new UpdateScript() {
+            final UpdateSummary summary = dc.executeUpdate(new UpdateScript() {
                 @Override
                 public void run(UpdateCallback cb) {
                     ColumnCreationBuilder createTableBuilder = 
cb.createTable(schema, "test_table").withColumn("id")
@@ -86,6 +87,7 @@ public class JdbcTestTemplates {
                     cb.insertInto(table).value("id", 4.0).value("code", 
"C02").execute();
                 }
             });
+            assertEquals(4, summary.getInsertedRows().get().intValue());
 
             assertEquals(1, 
getCount(dc.query().from("test_table").selectCount().where("code").isNull().execute()));
             assertEquals(3, 
getCount(dc.query().from("test_table").selectCount().where("code").isNotNull().execute()));

http://git-wip-us.apache.org/repos/asf/metamodel/blob/d888907b/jdbc/src/test/java/org/apache/metamodel/jdbc/integrationtests/PostgresqlTest.java
----------------------------------------------------------------------
diff --git 
a/jdbc/src/test/java/org/apache/metamodel/jdbc/integrationtests/PostgresqlTest.java
 
b/jdbc/src/test/java/org/apache/metamodel/jdbc/integrationtests/PostgresqlTest.java
index 5cf6822..ef3ae28 100644
--- 
a/jdbc/src/test/java/org/apache/metamodel/jdbc/integrationtests/PostgresqlTest.java
+++ 
b/jdbc/src/test/java/org/apache/metamodel/jdbc/integrationtests/PostgresqlTest.java
@@ -23,6 +23,7 @@ import java.sql.Connection;
 import java.sql.DatabaseMetaData;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 import javax.swing.table.TableModel;
@@ -31,8 +32,10 @@ import org.apache.metamodel.BatchUpdateScript;
 import org.apache.metamodel.DataContext;
 import org.apache.metamodel.UpdateCallback;
 import org.apache.metamodel.UpdateScript;
+import org.apache.metamodel.UpdateSummary;
 import org.apache.metamodel.data.DataSet;
 import org.apache.metamodel.data.DataSetTableModel;
+import org.apache.metamodel.drop.DropTable;
 import org.apache.metamodel.insert.RowInsertionBuilder;
 import org.apache.metamodel.jdbc.JdbcDataContext;
 import org.apache.metamodel.jdbc.JdbcTestTemplates;
@@ -430,6 +433,42 @@ public class PostgresqlTest extends 
AbstractJdbIntegrationTest {
             });
         }
     }
+
+    public void testGetGeneratedKeys() throws Exception {
+        if (!isConfigured()) {
+            return;
+        }
+
+        final JdbcDataContext dc = new JdbcDataContext(getConnection());
+        final Schema schema = dc.getDefaultSchema();
+        final String tableName = "my_table_with_generated_keys";
+        
+        if (schema.getTableByName(tableName) != null) {
+            dc.executeUpdate(new DropTable(schema, tableName));
+        }
+
+        final UpdateSummary updateSummary = dc.executeUpdate(new 
UpdateScript() {
+            @Override
+            public void run(UpdateCallback cb) {
+                Table table = cb.createTable(schema, 
tableName).withColumn("id").ofType(ColumnType.INTEGER)
+                        
.ofNativeType("SERIAL").nullable(false).asPrimaryKey().withColumn("foo").ofType(
+                                ColumnType.STRING).execute();
+                assertEquals(tableName, table.getName());
+
+                cb.insertInto(table).value("foo", "hello").execute();
+                cb.insertInto(table).value("foo", "world").execute();
+            }
+        });
+
+        final Optional<Integer> insertedRows = updateSummary.getInsertedRows();
+        assertTrue(insertedRows.isPresent());
+        assertEquals(2, insertedRows.get().intValue());
+        
+        final Optional<Iterable<Object>> generatedKeys = 
updateSummary.getGeneratedKeys();
+        assertTrue(generatedKeys.isPresent());
+        assertEquals("[1, 2]", generatedKeys.get().toString());
+        
+    }
     
     public void testBlob() throws Exception {
         if (!isConfigured()) {

Reply via email to