[core] better support for generated keys

JDBC offers an API to specify what keys to return. Rails knows the keys. Use that API directly to avoid loading too much data. E.g. pgjdbc appends "RETURNING *" to an insert query, effectively return _all_ columns. Since the jdbc driver does that, nobody ever noticed that the PG adapter is missing the code that appends the `RETURNING <pk>`
jruby · Dec 3, 2019 · c68df3f · c68df3f
1 parent bc55874
commit c68df3f
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 9 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,6 +2,7 @@
 *.tgz
 *~
 *.log
+/log/
 patches*
 *#
 TAGS
@@ -30,4 +31,4 @@ Gemfile.lock
 .idea
 .settings
 activerecord-jdbc.iml
-lib/arjdbc/jdbc/adapter_java.jar
+lib/arjdbc/jdbc/adapter_java.jar
diff --git a/lib/arjdbc/abstract/database_statements.rb b/lib/arjdbc/abstract/database_statements.rb
@@ -13,10 +13,10 @@ def exec_insert(sql, name = nil, binds = NO_BINDS, pk = nil, sequence_name = nil
         binds = convert_legacy_binds_to_attributes(binds) if binds.first.is_a?(Array)
 
         if without_prepared_statement?(binds)
-          log(sql, name) { @connection.execute_insert(sql) }
+          log(sql, name) { @connection.execute_insert(sql, pk) }
         else
           log(sql, name, binds) do
-            @connection.execute_insert(sql, binds)
+            @connection.execute_insert(sql, binds, pk)
           end
         end
       end

diff --git a/src/java/arjdbc/jdbc/RubyJdbcConnection.java b/src/java/arjdbc/jdbc/RubyJdbcConnection.java
@@ -888,23 +888,45 @@ protected IRubyObject mapExecuteResult(final ThreadContext context,
         return mapQueryResult(context, connection, resultSet);
     }
 
+    private static String[] createStatementPk(IRubyObject pk) {
+        String[] statementPk;
+        if (pk instanceof RubyArray) {
+            RubyArray ary = (RubyArray) pk;
+            int size = ary.size();
+            statementPk = new String[size];
+            for (int i = 0; i < size; i++) {
+                statementPk[i] = sqlString(ary.eltInternal(i));
+            }
+        } else {
+            statementPk = new String[] { sqlString(pk) };
+        }
+        return statementPk;
+    }
+
     /**
      * Executes an INSERT SQL statement
      * @param context
      * @param sql
+     * @param pk Rails PK
      * @return ActiveRecord::Result
      * @throws SQLException
      */
-    @JRubyMethod(name = "execute_insert", required = 1)
-    public IRubyObject execute_insert(final ThreadContext context, final IRubyObject sql) {
+    @JRubyMethod(name = "execute_insert", required = 2)
+    public IRubyObject execute_insert(final ThreadContext context, final IRubyObject sql, final IRubyObject pk) {
         return withConnection(context, new Callable<IRubyObject>() {
             public IRubyObject call(final Connection connection) throws SQLException {
                 Statement statement = null;
                 final String query = sqlString(sql);
                 try {
 
                     statement = createStatement(context, connection);
-                    statement.executeUpdate(query, Statement.RETURN_GENERATED_KEYS);
+
+                    if (pk == context.nil || pk == context.fals || !supportsGeneratedKeys(connection)) {
+                        statement.executeUpdate(query, Statement.RETURN_GENERATED_KEYS);
+                    } else {
+                        statement.executeUpdate(query, createStatementPk(pk));
+                    }
+
                     return mapGeneratedKeys(context, connection, statement);
 
                 } catch (final SQLException e) {
@@ -922,18 +944,24 @@ public IRubyObject call(final Connection connection) throws SQLException {
      * @param context
      * @param sql
      * @param binds RubyArray of values to be bound to the query
+     * @param pk Rails PK
      * @return ActiveRecord::Result
      * @throws SQLException
      */
-    @JRubyMethod(name = "execute_insert", required = 2)
-    public IRubyObject execute_insert(final ThreadContext context, final IRubyObject sql, final IRubyObject binds) {
+    @JRubyMethod(name = "execute_insert", required = 3)
+    public IRubyObject execute_insert(final ThreadContext context, final IRubyObject sql, final IRubyObject binds,
+                                      final IRubyObject pk) {
         return withConnection(context, new Callable<IRubyObject>() {
             public IRubyObject call(final Connection connection) throws SQLException {
                 PreparedStatement statement = null;
                 final String query = sqlString(sql);
                 try {
+                    if (pk == context.nil || pk == context.fals || !supportsGeneratedKeys(connection)) {
+                        statement = connection.prepareStatement(query, Statement.RETURN_GENERATED_KEYS);
+                    } else {
+                        statement = connection.prepareStatement(query, createStatementPk(pk));
+                    }
 
-                    statement = connection.prepareStatement(query, Statement.RETURN_GENERATED_KEYS);
                     setStatementParameters(context, connection, statement, (RubyArray) binds);
                     statement.executeUpdate();
                     return mapGeneratedKeys(context, connection, statement);