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/src/main/java/io/helidon/dbclient/jdbc/JdbcStatement.java

@@ -79,6 +79,15 @@ 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}.
  *
@@ -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);
  }
  }
 }