Issue helidon-io#2279 - Return generated IDs from INSERT statement.

Signed-off-by: Tomáš Kraus <[email protected]>

dbclient/dbclient/src/main/java/io/helidon/dbclient/DbResultDml.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2024 Oracle and/or its affiliates.
+ *
+ * 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 io.helidon.dbclient;
+
+import java.util.stream.Stream;
+
+/**
+ * Result of DML statement execution.
+ */
+public interface DbResultDml extends AutoCloseable {
+
+    /**
+     * Retrieves any auto-generated keys created as a result of executing this DML statement.
+     *
+     * @return the auto-generated keys
+     */
+    Stream<DbRow> generatedKeys();
+
+    /**
+     * Retrieve statement execution result.
+     *
+     * @return row count for Data Manipulation Language (DML) statements or {@code 0}
+     *         for statements that return nothing.
+     */
+    long result();
+
+    /**
+     * Create new instance of DML statement execution result.
+     *
+     * @param generatedKeys the auto-generated keys
+     * @param result the statement execution result
+     * @return new instance of DML statement execution result
+     */
+    static DbResultDml create(Stream<DbRow> generatedKeys, long result) {
+        return new DbResultDmlImpl(generatedKeys, result);
+    }
+
+}

dbclient/dbclient/src/main/java/io/helidon/dbclient/DbResultDmlImpl.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2024 Oracle and/or its affiliates.
+ *
+ * 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 io.helidon.dbclient;
+
+import java.util.Objects;
+import java.util.stream.Stream;
+
+// DbResultDml implementation
+record DbResultDmlImpl(Stream<DbRow> generatedKeys, long result) implements DbResultDml {
+
+    DbResultDmlImpl {
+        Objects.requireNonNull(generatedKeys, "List of auto-generated keys value is null");
+        if (result < 0) {
+            throw new IllegalArgumentException("Statement execution result value is less than 0");
+        }
+    }
+
+    @Override
+    public void close() throws Exception {
+        generatedKeys.close();
+    }
+
+}

dbclient/dbclient/src/main/java/io/helidon/dbclient/DbStatementDml.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2019, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,8 @@
  */
 package io.helidon.dbclient;
 
+import java.util.List;
+
 /**
  * Data Manipulation Language (DML) database statement.
  * A DML statement modifies records in the database and returns the number of modified records.
@@ -24,7 +26,35 @@ public interface DbStatementDml extends DbStatement<DbStatementDml> {
     /**
      * Execute this statement using the parameters configured with {@code params} and {@code addParams} methods.
      *
-     * @return The result of this statement.
+     * @return the result of this statement
      */
     long execute();
+
+    /**
+     * Execute {@code INSERT} statement using the parameters configured with {@code params} and {@code addParams} methods
+     * and return compound result with generated keys.
+     *
+     * @return the result of this statement with generated keys
+     */
+    DbResultDml insert();
+
+    /**
+     * Set auto-generated keys to be returned from the statement execution using {@link #insert()}.
+     * Only one method from {@link #returnGeneratedKeys()} and {@link #returnColumns(List)} may be used.
+     * This feature is database provider specific and some databases require specific columns to be set.
+     *
+     * @return updated db statement
+     */
+    DbStatementDml returnGeneratedKeys();
+
+    /**
+     * Set column names to be returned from the inserted row or rows from the statement execution using {@link #insert()}.
+     * Only one method from {@link #returnGeneratedKeys()} and {@link #returnColumns(List)} may be used.
+     * This feature is database provider specific.
+     *
+     * @param columnNames an array of column names indicating the columns that should be returned from the inserted row or rows
+     * @return updated db statement
+     */
+    DbStatementDml returnColumns(List<String> columnNames);
+
 }

dbclient/jdbc/etc/spotbugs/exclude.xml

@@ -26,4 +26,16 @@
         <Method name="prepareStatement"/>
         <Bug pattern="SQL_INJECTION_JDBC"/>
     </Match>
+    <Match>
+        <!-- Doesn't construct SQL string. It converts string to statement -->
+        <Class name="io.helidon.dbclient.jdbc.JdbcStatementDml"/>
+        <Method name="prepareStatement"/>
+        <Bug pattern="SQL_INJECTION_JDBC"/>
+    </Match>
+    <Match>
+        <!-- Doesn't construct SQL string. It converts string to statement -->
+        <Class name="io.helidon.dbclient.jdbc.JdbcTransactionStatementDml"/>
+        <Method name="prepareStatement"/>
+        <Bug pattern="SQL_INJECTION_JDBC"/>
+    </Match>
 </FindBugsFilter>

dbclient/jdbc/src/main/java/io/helidon/dbclient/jdbc/JdbcStatement.java

@@ -79,11 +79,20 @@ private JdbcExecuteContext jdbcContext() {
         return context(JdbcExecuteContext.class);
     }
 
+    /**
+     * Set the connection.
+     *
+     * @param connection the database connection
+     */
+    protected void connection(Connection connection) {
+        this.connection = connection;
+    }
+
     /**
      * Create the {@link PreparedStatement}.
      *
      * @param serviceContext client service context
-     * @return PreparedStatement
+     * @return new instance of {@link PreparedStatement}
      */
     protected PreparedStatement prepareStatement(DbClientServiceContext serviceContext) {
         String stmtName = serviceContext.statementName();
@@ -105,7 +114,7 @@ protected PreparedStatement prepareStatement(DbClientServiceContext serviceConte
      *
      * @param stmtName statement name
      * @param stmt     statement text
-     * @return statement
+     * @return new instance of {@link PreparedStatement}
      */
     protected PreparedStatement prepareStatement(String stmtName, String stmt) {
         Connection connection = connectionPool.connection();
@@ -120,10 +129,10 @@ protected PreparedStatement prepareStatement(String stmtName, String stmt) {
     /**
      * Create the {@link PreparedStatement}.
      *
-     * @param connection connection
+     * @param connection the database connection
      * @param stmtName   statement name
      * @param stmt       statement text
-     * @return statement
+     * @return new instance of {@link PreparedStatement}
      */
     protected PreparedStatement prepareStatement(Connection connection, String stmtName, String stmt) {
         try {

dbclient/jdbc/src/main/java/io/helidon/dbclient/jdbc/JdbcStatementDml.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2019, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,11 +15,22 @@
  */
 package io.helidon.dbclient.jdbc;
 
+import java.sql.Connection;
 import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
+import io.helidon.dbclient.DbClientException;
 import io.helidon.dbclient.DbClientServiceContext;
+import io.helidon.dbclient.DbResultDml;
+import io.helidon.dbclient.DbRow;
 import io.helidon.dbclient.DbStatementDml;
 import io.helidon.dbclient.DbStatementException;
 import io.helidon.dbclient.DbStatementType;
@@ -29,7 +40,17 @@
  */
 class JdbcStatementDml extends JdbcStatement<DbStatementDml> implements DbStatementDml {
 
+    static final String[] EMPTY_STRING_ARRAY = new String[0];
+
     private final DbStatementType type;
+    // Column names to be returned from the inserted row or rows from the statement execution.
+    // Value of null (default) indicates no columns are set.
+    private List<String> columnNames = List.of();
+    // Whether PreparedStatement shall be created with Statement.RETURN_GENERATED_KEYS:
+    // - value of false (default) indicates that autoGeneratedKeys won't be passed to PreparedStatement creation
+    // - value of true indicates that Statement.RETURN_GENERATED_KEYS as autoGeneratedKeys will be passed
+    //   to PreparedStatement creation
+    private boolean returnGeneratedKeys;
 
     /**
      * Create a new instance.
@@ -58,6 +79,45 @@ public long execute() {
         });
     }
 
+    @Override
+    public DbResultDml insert() {
+        return doExecute((future, context) -> doInsert(this, future, context, this::closeConnection));
+    }
+
+    @Override
+    public DbStatementDml returnGeneratedKeys() {
+        if (!columnNames.isEmpty()) {
+            throw new IllegalStateException("Method returnColumns(String[]) was already called to set specific column names.");
+        }
+        returnGeneratedKeys = true;
+        return this;
+    }
+
+    @Override
+    public DbStatementDml returnColumns(List<String> columnNames) {
+        if (returnGeneratedKeys) {
+            throw new IllegalStateException("Method returnGeneratedKeys() was already called.");
+        }
+        Objects.requireNonNull(columnNames, "List of column names value is null");
+        this.columnNames = Collections.unmodifiableList(columnNames);
+        return this;
+    }
+
+    @Override
+    protected PreparedStatement prepareStatement(Connection connection, String stmtName, String stmt) {
+        try {
+            connection(connection);
+            if (returnGeneratedKeys) {
+                return connection.prepareStatement(stmt, Statement.RETURN_GENERATED_KEYS);
+            } else if (!columnNames.isEmpty()) {
+                return connection.prepareStatement(stmt, columnNames.toArray(EMPTY_STRING_ARRAY));
+            }
+            return connection.prepareStatement(stmt);
+        } catch (SQLException e) {
+            throw new DbClientException(String.format("Failed to prepare statement: %s", stmtName), e);
+        }
+    }
+
     /**
      * Execute the given statement.
      *
@@ -75,7 +135,41 @@ static long doExecute(JdbcStatement<? extends DbStatementDml> dbStmt,
             future.complete(result);
             return result;
         } catch (SQLException ex) {
+            dbStmt.closeConnection();
             throw new DbStatementException("Failed to execute statement", dbStmt.context().statement(), ex);
         }
     }
+
+    /**
+     * Execute the given insert statement.
+     *
+     * @param dbStmt  db statement
+     * @param future  query future
+     * @param context service context
+     * @return query result
+     */
+    static DbResultDml doInsert(JdbcStatement<? extends DbStatementDml> dbStmt,
+                                CompletableFuture<Long> future,
+                                DbClientServiceContext context,
+                                Runnable onClose) {
+        PreparedStatement statement;
+        try {
+            statement = dbStmt.prepareStatement(context);
+            long result = statement.executeUpdate();
+            ResultSet rs = statement.getGeneratedKeys();
+            JdbcRow.Spliterator spliterator = new JdbcRow.Spliterator(rs, statement, dbStmt.context(), future);
+            Stream<DbRow> generatedKeys = autoClose(StreamSupport.stream(spliterator, false)
+                    .onClose(() -> {
+                        spliterator.close();
+                        if (onClose != null) {
+                            onClose.run();
+                        }
+                    }));
+            return DbResultDml.create(generatedKeys, result);
+        } catch (SQLException ex) {
+            dbStmt.closeConnection();
+            throw new DbStatementException("Failed to execute statement", dbStmt.context().statement(), ex);
+        }
+    }
+
 }

dbclient/jdbc/src/main/java/io/helidon/dbclient/jdbc/JdbcStatementQuery.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2019, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -82,7 +82,7 @@ static Stream<DbRow> doExecute(JdbcStatement<? extends DbStatementQuery> dbStmt,
                     }));
         } catch (SQLException ex) {
             dbStmt.closeConnection();
-            throw new DbStatementException("Failed to create Statement", dbStmt.context().statement(), ex);
+            throw new DbStatementException("Failed to execute Statement", dbStmt.context().statement(), ex);
         }
     }
 }