refactor: deduplicate shared logic and remove dead code across modules

storm-core/src/main/java/st/orm/core/repository/impl/EntityRepositoryImpl.java

@@ -30,6 +30,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -1717,16 +1718,7 @@ protected void insert(List<E> batch, PreparedQuery query, boolean ignoreAutoGene
         if (batch.isEmpty()) {
             return;
         }
-        List<E> transformed = batch.stream()
-                .map(this::fireBeforeInsert)
-                .toList();
-        transformed.stream()
-                .map(e -> validateInsert(e, ignoreAutoGenerate))
-                .forEach(query::addBatch);
-        int[] result = query.executeBatch();
-        if (IntStream.of(result).anyMatch(r -> r != 1)) {
-            throw new PersistenceException("Batch insert of %s failed. One or more rows were not affected.".formatted(model.type().getSimpleName()));
-        }
+        List<E> transformed = executeInsertBatch(batch, query, ignoreAutoGenerate);
         transformed.forEach(this::fireAfterInsert);
     }
 
@@ -1739,16 +1731,7 @@ private List<ID> insertAndFetchIds(List<E> batch, PreparedQuery query, boolean i
         if (batch.isEmpty()) {
             return List.of();
         }
-        List<E> transformed = batch.stream()
-                .map(this::fireBeforeInsert)
-                .toList();
-        transformed.stream()
-                .map(e -> validateInsert(e, ignoreAutoGenerate))
-                .forEach(query::addBatch);
-        int[] result = query.executeBatch();
-        if (IntStream.of(result).anyMatch(r -> r != 1)) {
-            throw new PersistenceException("Batch insert of %s failed. One or more rows were not affected.".formatted(model.type().getSimpleName()));
-        }
+        List<E> transformed = executeInsertBatch(batch, query, ignoreAutoGenerate);
         List<ID> ids;
         if (isAutoGeneratedPrimaryKey() && !ignoreAutoGenerate) {
             try (var stream = query.getGeneratedKeys(model.primaryKeyType())) {
@@ -1761,6 +1744,26 @@ private List<ID> insertAndFetchIds(List<E> batch, PreparedQuery query, boolean i
         return ids;
     }
 
+    /**
+     * Runs the batch through the prepared insert: fires the before-insert callbacks, validates, executes the
+     * batch and checks that every row was inserted.
+     *
+     * @return the entities as transformed by the before-insert callbacks.
+     */
+    private List<E> executeInsertBatch(List<E> batch, PreparedQuery query, boolean ignoreAutoGenerate) {
+        List<E> transformed = batch.stream()
+                .map(this::fireBeforeInsert)
+                .toList();
+        transformed.stream()
+                .map(e -> validateInsert(e, ignoreAutoGenerate))
+                .forEach(query::addBatch);
+        int[] result = query.executeBatch();
+        if (IntStream.of(result).anyMatch(r -> r != 1)) {
+            throw new PersistenceException("Batch insert of %s failed. One or more rows were not affected.".formatted(model.type().getSimpleName()));
+        }
+        return transformed;
+    }
+
     /**
      * Updates a stream of entities in the database using the default batch size.
      *
@@ -1984,22 +1987,7 @@ protected PreparedQuery prepareUpdateQuery(Set<Metamodel<?, ?>> fields) {
     }
 
     protected void update(List<E> batch, PreparedQuery query, @Nullable EntityCache<E, ID> cache) {
-        if (batch.isEmpty()) {
-            return;
-        }
-        batch.stream().map(this::validateUpdate).forEach(query::addBatch);
-        if (cache != null) {
-            batch.stream()
-                    .filter(e -> !model.isDefaultPrimaryKey(e.id()))
-                    .forEach(e -> cache.remove(e.id()));
-        }
-        int[] result = query.executeBatch();
-        if (query.isVersionAware() && IntStream.of(result).anyMatch(r -> r == 0)) {
-            throw new OptimisticLockException("Batch update of %s failed due to optimistic lock. One or more entities may have been modified or deleted by another transaction.".formatted(model.type().getSimpleName()));
-        } else if (IntStream.of(result).anyMatch(r -> r != 1)) {
-            throw new PersistenceException("Batch update of %s failed. One or more rows were not affected.".formatted(model.type().getSimpleName()));
-        }
-        batch.forEach(this::fireAfterUpdate);
+        updateAndFetchIds(batch, query, cache);
     }
 
     protected List<ID> updateAndFetchIds(List<E> batch, PreparedQuery query, @Nullable EntityCache<E, ID> cache) {
@@ -2168,6 +2156,124 @@ protected List<ID> doUpsertAndFetchIdsBatch(List<E> batch, PreparedQuery query,
         throw upsertNotAvailable();
     }
 
+    // Partition keys for the dialects' id-returning upsertAndFetchIds implementations.
+    public sealed interface SeqPartitionKey {}
+    public static final class SeqNoOpKey implements SeqPartitionKey {
+        public static final SeqNoOpKey INSTANCE = new SeqNoOpKey();
+        private SeqNoOpKey() {
+        }
+    }
+    public static final class SeqUpsertKey implements SeqPartitionKey {
+        public static final SeqUpsertKey INSTANCE = new SeqUpsertKey();
+        private SeqUpsertKey() {
+        }
+    }
+    public record SeqUpdateKey(Set<Metamodel<?, ?>> fields) implements SeqPartitionKey {
+        public SeqUpdateKey() {
+            this(Set.of()); // All fields.
+        }
+    }
+
+    /**
+     * The upsert-and-fetch-ids loop for dialects that fetch the ids through a single id-returning query
+     * (such as {@code INSERT ... RETURNING} or {@code MERGE ... OUTPUT}) rather than batched prepared
+     * statements: the entities are partitioned into no-op, upsert and per-shape update batches, and the
+     * resulting ids come back in entity order.
+     *
+     * @param entities the entities to upsert.
+     * @param upsertPartition executes one upsert partition and returns its ids; receives the partition's
+     *                        chunk and the entity cache in effect for the operation. Dialects with the
+     *                        standard cache-eviction semantics pass {@link #upsertPartitionAndFetchIds}.
+     * @return the ids of the upserted entities.
+     * @since 1.14
+     */
+    protected List<ID> upsertAndFetchIdsPartitioned(
+            Iterable<E> entities,
+            BiFunction<List<E>, Optional<EntityCache<E, ID>>, List<ID>> upsertPartition) {
+        var updateQueries = new HashMap<Set<Metamodel<?, ?>>, PreparedQuery>();
+        try {
+            var result = new ArrayList<ID>();
+            var entityCache = entityCache();
+            partitioned(toStream(entities), defaultBatchSize, entity -> {
+                if (isUpsertUpdate(entity)) {
+                    var dirty = getDirty(entity, entityCache.orElse(null));
+                    if (dirty.isEmpty()) {
+                        return SeqNoOpKey.INSTANCE;
+                    }
+                    return new SeqUpdateKey(dirty.get());
+                } else {
+                    return SeqUpsertKey.INSTANCE;
+                }
+            }, getMaxShapes(), new SeqUpdateKey()).forEach(partition -> {
+                switch (partition.key()) {
+                    case SeqNoOpKey ignore -> result.addAll(partition.chunk().stream().map(E::id).toList());
+                    case SeqUpsertKey ignore -> result.addAll(upsertPartition.apply(partition.chunk(), entityCache));
+                    case SeqUpdateKey u -> {
+                        List<E> batch = hasEntityCallbacks()
+                                ? partition.chunk().stream().map(this::fireBeforeUpdate).toList()
+                                : partition.chunk();
+                        result.addAll(updateAndFetchIds(batch,
+                                updateQueries.computeIfAbsent(u.fields(), this::prepareUpdateQuery),
+                                entityCache.orElse(null)));
+                    }
+                }
+            });
+            return result;
+        } finally {
+            closeQuietly(updateQueries.values().stream());
+        }
+    }
+
+    /**
+     * Executes one upsert partition through a single id-returning query: fires the before-upsert callbacks,
+     * evicts the entities with non-default primary keys from the cache (the upsert may update them), runs
+     * the query and reports the returned ids to the after-upsert callbacks.
+     *
+     * @param chunk the partition's entities.
+     * @param entityCache the entity cache in effect for the operation.
+     * @param upsertQuery creates the id-returning upsert query for a batch.
+     * @return the ids of the upserted entities.
+     * @since 1.14
+     */
+    protected List<ID> upsertPartitionAndFetchIds(List<E> chunk,
+                                                  Optional<EntityCache<E, ID>> entityCache,
+                                                  Function<List<E>, Query> upsertQuery) {
+        List<E> batch = hasEntityCallbacks()
+                ? chunk.stream().map(this::fireBeforeUpsert).toList()
+                : chunk;
+        entityCache.ifPresent(cache -> batch.stream()
+                .filter(e -> !model.isDefaultPrimaryKey(e.id()))
+                .forEach(e -> cache.remove(e.id())));
+        List<ID> ids = upsertQuery.apply(batch).getResultList(model.primaryKeyType());
+        fireAfterUpsert(batch, ids);
+        return ids;
+    }
+
+    /**
+     * Inserts the entity and fetches the generated primary key through the dialect's
+     * {@code INSERT ... RETURNING} clause, for dialects whose driver cannot report sequence-generated
+     * keys through {@code getGeneratedKeys()}.
+     *
+     * @param entity the entity to insert.
+     * @return the primary key of the inserted entity.
+     * @since 1.14
+     */
+    protected ID insertAndFetchIdReturning(E entity) {
+        entity = fireBeforeInsert(entity);
+        validateInsert(entity);
+        assert primaryKeyColumns.size() == 1;
+        var primaryKeyColumn = primaryKeyColumns.getFirst();
+        String pkName = primaryKeyColumn.qualifiedName(ormTemplate.dialect());
+        try (var query = ormTemplate.query(TemplateString.raw("""
+                INSERT INTO \0
+                VALUES \0
+                RETURNING %s""".formatted(pkName), model.type(), entity)).managed().prepare()) {
+            ID id = query.getSingleResult(model.primaryKeyType());
+            fireAfterInsert(entity, id);
+            return id;
+        }
+    }
+
     /**
      * Removes a stream of entities from the database in batches.
      *

storm-core/src/main/java/st/orm/core/repository/impl/MergeEntityRepositoryImpl.java

@@ -278,17 +278,7 @@ private TemplateString mergeStatement(TemplateString mergeSelect, AtomicBoolean
      */
     @Override
     protected void doUpsert(E entity) {
-        validateUpsert(entity);
-        entityCache().ifPresent(cache -> {
-            if (!model.isDefaultPrimaryKey(entity.id())) {
-                cache.remove(entity.id());
-            }
-        });
-        var versionAware = new AtomicBoolean();
-        intercept(sql -> sql.versionAware(versionAware.getPlain()), () -> {
-            var query = ormTemplate.query(mergeStatement(mergeSelect(entity), versionAware)).managed();
-            query.executeUpdate();
-        });
+        doUpsertAndFetchId(entity);
     }
 
     /**
@@ -322,19 +312,7 @@ protected PreparedQuery prepareUpsertQuery() {
     @Override
     protected void doUpsertBatch(List<E> batch, PreparedQuery query,
                                  @Nullable EntityCache<E, ID> cache) {
-        if (batch.isEmpty()) {
-            return;
-        }
-        batch.stream().map(this::validateUpsert).forEach(query::addBatch);
-        if (cache != null) {
-            batch.stream()
-                    .filter(e -> !model.isDefaultPrimaryKey(e.id()))
-                    .forEach(e -> cache.remove(e.id()));
-        }
-        int[] result = query.executeBatch();
-        if (IntStream.of(result).anyMatch(r -> r != 0 && r != 1 && r != 2)) {
-            throw new PersistenceException(upsertFailureMessage(batch.size()));
-        }
+        doUpsertAndFetchIdsBatch(batch, query, cache);
     }
 
     @Override