Add the openStream method to the Criteria API to return a Stream (#1111)

nakamura-to · web-flow · commit 7593bf42ed2e · 2024-05-26T21:42:16.000+09:00
diff --git a/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSelectIntermediate.java b/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSelectIntermediate.java
@@ -33,24 +33,31 @@ public SetOperationContext<ELEMENT> getContext() {
     return new SetOperationContext.Select<>(declaration.getContext());
   }
 
+  @Override
+  public Stream<ELEMENT> openStream() {
+    NativeSqlSelectTerminal<Stream<ELEMENT>> terminal =
+        createNativeSqlSelectTerminal(Function.identity(), true);
+    return terminal.execute();
+  }
+
   @Override
   public <RESULT> RESULT mapStream(Function<Stream<ELEMENT>, RESULT> streamMapper) {
     Objects.requireNonNull(streamMapper);
-    NativeSqlSelectTerminal<RESULT> terminal = createNativeSqlSelectTerminal(streamMapper);
+    NativeSqlSelectTerminal<RESULT> terminal = createNativeSqlSelectTerminal(streamMapper, false);
     return terminal.execute();
   }
 
   @Override
   protected Command<List<ELEMENT>> createCommand() {
     NativeSqlSelectTerminal<List<ELEMENT>> terminal =
-        createNativeSqlSelectTerminal(stream -> stream.collect(toList()));
+        createNativeSqlSelectTerminal(stream -> stream.collect(toList()), false);
     return terminal.createCommand();
   }
 
   private <RESULT> NativeSqlSelectTerminal<RESULT> createNativeSqlSelectTerminal(
-      Function<Stream<ELEMENT>, RESULT> streamMapper) {
+      Function<Stream<ELEMENT>, RESULT> streamMapper, boolean returnsStream) {
     ResultSetHandler<RESULT> handler =
         new MappedResultStreamHandler<>(streamMapper, objectProviderFactory);
-    return new NativeSqlSelectTerminal<>(config, declaration, handler);
+    return new NativeSqlSelectTerminal<>(config, declaration, handler, returnsStream);
   }
 }
diff --git a/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSelectStarting.java b/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSelectStarting.java
@@ -838,6 +838,13 @@ public SetOperationContext<ENTITY> getContext() {
     return intermediate.getContext();
   }
 
+  @Override
+  public Stream<ENTITY> openStream() {
+    NativeSqlSelectIntermediate<ENTITY> intermediate =
+        new NativeSqlSelectIntermediate<>(config, declaration, createEntityProviderFactory());
+    return intermediate.openStream();
+  }
+
   @Override
   public <RESULT> RESULT mapStream(Function<Stream<ENTITY>, RESULT> streamMapper) {
     NativeSqlSelectIntermediate<ENTITY> intermediate =
diff --git a/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSelectTerminal.java b/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSelectTerminal.java
@@ -1,6 +1,7 @@
 package org.seasar.doma.jdbc.criteria.statement;
 
 import java.util.Objects;
+import org.seasar.doma.FetchType;
 import org.seasar.doma.jdbc.Config;
 import org.seasar.doma.jdbc.PreparedSql;
 import org.seasar.doma.jdbc.command.Command;
@@ -18,11 +19,22 @@ public class NativeSqlSelectTerminal<RESULT>
   private final SelectFromDeclaration declaration;
   private final ResultSetHandler<RESULT> resultSetHandler;
 
+  private final Boolean returnsStream;
+
   public NativeSqlSelectTerminal(
       Config config, SelectFromDeclaration declaration, ResultSetHandler<RESULT> resultSetHandler) {
+    this(config, declaration, resultSetHandler, false);
+  }
+
+  public NativeSqlSelectTerminal(
+      Config config,
+      SelectFromDeclaration declaration,
+      ResultSetHandler<RESULT> resultSetHandler,
+      boolean returnsStream) {
     super(Objects.requireNonNull(config));
     this.declaration = Objects.requireNonNull(declaration);
     this.resultSetHandler = Objects.requireNonNull(resultSetHandler);
+    this.returnsStream = returnsStream;
   }
 
   /**
@@ -46,10 +58,7 @@ protected Command<RESULT> createCommand() {
         new SelectBuilder(
             config, context, createCommenter(settings.getComment()), settings.getSqlLogType());
     PreparedSql sql = builder.build();
-    CriteriaQuery query = new CriteriaQuery(config, sql, getClass().getName(), EXECUTE_METHOD_NAME);
-    query.setFetchSize(settings.getFetchSize());
-    query.setMaxRows(settings.getMaxRows());
-    query.setQueryTimeout(settings.getQueryTimeout());
+    CriteriaQuery query = createCriteriaQuery(sql, settings);
     return new SelectCommand<RESULT>(query, resultSetHandler) {
       @Override
       public RESULT execute() {
@@ -62,4 +71,24 @@ public RESULT execute() {
       }
     };
   }
+
+  private CriteriaQuery createCriteriaQuery(PreparedSql sql, SelectSettings settings) {
+    CriteriaQuery query =
+        new CriteriaQuery(config, sql, getClass().getName(), EXECUTE_METHOD_NAME) {
+
+          @Override
+          public boolean isResultStream() {
+            return returnsStream;
+          }
+
+          @Override
+          public FetchType getFetchType() {
+            return returnsStream ? FetchType.LAZY : FetchType.EAGER;
+          }
+        };
+    query.setFetchSize(settings.getFetchSize());
+    query.setMaxRows(settings.getMaxRows());
+    query.setQueryTimeout(settings.getQueryTimeout());
+    return query;
+  }
 }
diff --git a/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSetStarting.java b/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSetStarting.java
@@ -43,24 +43,31 @@ public SetOperationContext<ELEMENT> getContext() {
     return context;
   }
 
+  @Override
+  public Stream<ELEMENT> openStream() {
+    NativeSqlSetTerminal<Stream<ELEMENT>> terminal =
+        createNativeSqlSetTerminal(Function.identity(), true);
+    return terminal.execute();
+  }
+
   @Override
   public <RESULT> RESULT mapStream(Function<Stream<ELEMENT>, RESULT> streamMapper) {
     Objects.requireNonNull(streamMapper);
-    NativeSqlSetTerminal<RESULT> terminal = createNativeSqlSetTerminal(streamMapper);
+    NativeSqlSetTerminal<RESULT> terminal = createNativeSqlSetTerminal(streamMapper, false);
     return terminal.execute();
   }
 
   @Override
   protected Command<List<ELEMENT>> createCommand() {
     NativeSqlSetTerminal<List<ELEMENT>> terminal =
-        createNativeSqlSetTerminal(stream -> stream.collect(toList()));
+        createNativeSqlSetTerminal(stream -> stream.collect(toList()), false);
     return terminal.createCommand();
   }
 
   private <RESULT> NativeSqlSetTerminal<RESULT> createNativeSqlSetTerminal(
-      Function<Stream<ELEMENT>, RESULT> streamMapper) {
+      Function<Stream<ELEMENT>, RESULT> streamMapper, boolean returnsStream) {
     ResultSetHandler<RESULT> handler =
         new MappedResultStreamHandler<>(streamMapper, objectProviderFactory);
-    return new NativeSqlSetTerminal<>(config, context, handler);
+    return new NativeSqlSetTerminal<>(config, context, handler, returnsStream);
   }
 }
diff --git a/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSetTerminal.java b/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/NativeSqlSetTerminal.java
@@ -1,6 +1,7 @@
 package org.seasar.doma.jdbc.criteria.statement;
 
 import java.util.Objects;
+import org.seasar.doma.FetchType;
 import org.seasar.doma.jdbc.Config;
 import org.seasar.doma.jdbc.PreparedSql;
 import org.seasar.doma.jdbc.command.Command;
@@ -17,11 +18,22 @@ public class NativeSqlSetTerminal<RESULT>
   private final SetOperationContext<?> context;
   private final ResultSetHandler<RESULT> resultSetHandler;
 
+  private final boolean returnsStream;
+
   public NativeSqlSetTerminal(
       Config config, SetOperationContext<?> context, ResultSetHandler<RESULT> resultSetHandler) {
+    this(config, context, resultSetHandler, false);
+  }
+
+  public NativeSqlSetTerminal(
+      Config config,
+      SetOperationContext<?> context,
+      ResultSetHandler<RESULT> resultSetHandler,
+      boolean returnsStream) {
     super(Objects.requireNonNull(config));
     this.context = Objects.requireNonNull(context);
     this.resultSetHandler = Objects.requireNonNull(resultSetHandler);
+    this.returnsStream = returnsStream;
   }
 
   /**
@@ -42,11 +54,28 @@ protected Command<RESULT> createCommand() {
         new SetOperationBuilder(
             config, context, createCommenter(settings.getComment()), settings.getSqlLogType());
     PreparedSql sql = builder.build();
-    CriteriaQuery query = new CriteriaQuery(config, sql, getClass().getName(), EXECUTE_METHOD_NAME);
+    CriteriaQuery query = createCriteriaQuery(sql, settings);
+    return new SelectCommand<>(query, resultSetHandler);
+  }
+
+  private CriteriaQuery createCriteriaQuery(PreparedSql sql, SelectSettings settings) {
+    CriteriaQuery query =
+        new CriteriaQuery(config, sql, getClass().getName(), EXECUTE_METHOD_NAME) {
+
+          @Override
+          public boolean isResultStream() {
+            return returnsStream;
+          }
+
+          @Override
+          public FetchType getFetchType() {
+            return returnsStream ? FetchType.LAZY : FetchType.EAGER;
+          }
+        };
     query.setFetchSize(settings.getFetchSize());
     query.setMaxRows(settings.getMaxRows());
     query.setQueryTimeout(settings.getQueryTimeout());
-    return new SelectCommand<>(query, resultSetHandler);
+    return query;
   }
 
   private SelectSettings findSettings() {
diff --git a/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/StreamMappable.java b/doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/StreamMappable.java
@@ -14,6 +14,15 @@
  */
 public interface StreamMappable<ELEMENT> extends Listable<ELEMENT> {
 
+  /**
+   * Open a stream.
+   *
+   * <p>You must close the stream after using it.
+   *
+   * @return the opened stream
+   */
+  Stream<ELEMENT> openStream();
+
   /**
    * Map a stream.
    *
diff --git a/doma-core/src/test/java/org/seasar/doma/jdbc/criteria/NativeSqlSelectTest.java b/doma-core/src/test/java/org/seasar/doma/jdbc/criteria/NativeSqlSelectTest.java
@@ -1,7 +1,9 @@
 package org.seasar.doma.jdbc.criteria;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.seasar.doma.jdbc.criteria.expression.Expressions.alias;
 import static org.seasar.doma.jdbc.criteria.expression.Expressions.avg;
 import static org.seasar.doma.jdbc.criteria.expression.Expressions.avgAsDouble;
@@ -24,6 +26,7 @@
 import java.math.BigDecimal;
 import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.seasar.doma.DomaException;
 import org.seasar.doma.internal.jdbc.mock.MockConfig;
@@ -57,7 +60,9 @@
 
 class NativeSqlSelectTest {
 
-  private final NativeSql nativeSql = new NativeSql(new MockConfig());
+  private final MockConfig config = new MockConfig();
+
+  private final NativeSql nativeSql = new NativeSql(config);
 
   @Test
   void from() {
@@ -1607,4 +1612,24 @@ private static UserDefinedExpression<Long> countDistinctMultiple(
           c.appendSql("))");
         });
   }
+
+  @Test
+  void openStream() {
+    Emp_ e = new Emp_();
+
+    Stream<Emp> stream = nativeSql.from(e).openStream();
+    assertFalse(config.dataSource.connection.closed);
+    stream.close();
+    assertTrue(config.dataSource.connection.closed);
+  }
+
+  @Test
+  void openStream_union() {
+    Emp_ e = new Emp_();
+
+    Stream<Emp> stream = nativeSql.from(e).union(nativeSql.from(e)).openStream();
+    assertFalse(config.dataSource.connection.closed);
+    stream.close();
+    assertTrue(config.dataSource.connection.closed);
+  }
 }
diff --git a/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KNativeSqlSelectIntermediate.kt b/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KNativeSqlSelectIntermediate.kt
@@ -3,13 +3,18 @@ package org.seasar.doma.kotlin.jdbc.criteria.statement
 import org.seasar.doma.jdbc.Sql
 import org.seasar.doma.jdbc.criteria.context.SetOperationContext
 import org.seasar.doma.jdbc.criteria.statement.SetOperand
+import java.util.stream.Stream
 import kotlin.streams.asSequence
 
 class KNativeSqlSelectIntermediate<ELEMENT>(private val statement: SetOperand<ELEMENT>) : KStatement<List<ELEMENT>>, KSetOperand<ELEMENT> {
 
     override val context: SetOperationContext<ELEMENT>
         get() = statement.context
 
+    override fun openStream(): Stream<ELEMENT> {
+        return statement.openStream()
+    }
+
     override fun <RESULT> mapSequence(sequenceMapper: (Sequence<ELEMENT>) -> RESULT): RESULT {
         return statement.mapStream {
             sequenceMapper(it.asSequence())
diff --git a/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KNativeSqlSelectStarting.kt b/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KNativeSqlSelectStarting.kt
@@ -26,6 +26,7 @@ import org.seasar.doma.kotlin.jdbc.criteria.declaration.KJoinDeclaration
 import org.seasar.doma.kotlin.jdbc.criteria.declaration.KOrderByNameDeclaration
 import org.seasar.doma.kotlin.jdbc.criteria.declaration.KWhereDeclaration
 import java.util.function.Function
+import java.util.stream.Stream
 import kotlin.streams.asSequence
 
 class KNativeSqlSelectStarting<ENTITY>(private val statement: NativeSqlSelectStarting<ENTITY>) :
@@ -401,6 +402,10 @@ class KNativeSqlSelectStarting<ENTITY>(private val statement: NativeSqlSelectSta
         return statement.execute()
     }
 
+    override fun openStream(): Stream<ENTITY> {
+        return statement.openStream()
+    }
+
     override fun <RESULT> mapSequence(sequenceMapper: (Sequence<ENTITY>) -> RESULT): RESULT {
         return statement.mapStream {
             sequenceMapper(it.asSequence())
diff --git a/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KNativeSqlSetStarting.kt b/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KNativeSqlSetStarting.kt
@@ -5,6 +5,7 @@ import org.seasar.doma.jdbc.criteria.context.SetOperationContext
 import org.seasar.doma.jdbc.criteria.statement.SetOperand
 import org.seasar.doma.jdbc.criteria.statement.SetOperator
 import org.seasar.doma.kotlin.jdbc.criteria.declaration.KOrderByIndexDeclaration
+import java.util.stream.Stream
 import kotlin.streams.asSequence
 
 class KNativeSqlSetStarting<ELEMENT>(val statement: SetOperator<ELEMENT>) : KSetOperator<ELEMENT> {
@@ -27,6 +28,17 @@ class KNativeSqlSetStarting<ELEMENT>(val statement: SetOperator<ELEMENT>) : KSet
         return statement.asSql()
     }
 
+    /**
+     * Open a stream.
+     *
+     * You must close the stream after using it.
+     *
+     * @return the opened stream
+     */
+    override fun openStream(): Stream<ELEMENT> {
+        return statement.openStream()
+    }
+
     override fun <RESULT> mapSequence(sequenceMapper: (Sequence<ELEMENT>) -> RESULT): RESULT {
         return statement.mapStream {
             sequenceMapper(it.asSequence())
diff --git a/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KSequenceMappable.kt b/doma-kotlin/src/main/kotlin/org/seasar/doma/kotlin/jdbc/criteria/statement/KSequenceMappable.kt
@@ -1,6 +1,10 @@
 package org.seasar.doma.kotlin.jdbc.criteria.statement
 
+import java.util.stream.Stream
+
 interface KSequenceMappable<ELEMENT> : KListable<ELEMENT> {
 
+    fun openStream(): Stream<ELEMENT>
+
     fun <RESULT> mapSequence(sequenceMapper: (Sequence<ELEMENT>) -> RESULT): RESULT
 }
diff --git a/doma-kotlin/src/test/kotlin/org/seasar/doma/jdbc/criteria/KNativeSqlSelectTest.kt b/doma-kotlin/src/test/kotlin/org/seasar/doma/jdbc/criteria/KNativeSqlSelectTest.kt
@@ -1,9 +1,11 @@
 package org.seasar.doma.jdbc.criteria
 
 import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertFalse
 import org.junit.jupiter.api.Assertions.assertThrows
 import org.junit.jupiter.api.Test
 import org.seasar.doma.DomaException
+import org.seasar.doma.internal.util.AssertionUtil.assertTrue
 import org.seasar.doma.jdbc.CommentContext
 import org.seasar.doma.jdbc.Commenter
 import org.seasar.doma.jdbc.criteria.entity.Dept_
@@ -45,7 +47,8 @@ import java.math.BigDecimal
 
 internal class KNativeSqlSelectTest {
 
-    private val nativeSql = org.seasar.doma.kotlin.jdbc.criteria.KNativeSql(MockConfig())
+    private val config = MockConfig()
+    private val nativeSql = org.seasar.doma.kotlin.jdbc.criteria.KNativeSql(config)
 
     @Test
     fun from() {
@@ -1770,4 +1773,15 @@ internal class KNativeSqlSelectTest {
             appendSql(" + 100)")
         }
     }
+
+    @Test
+    fun openStream() {
+        val e = Emp_()
+        val stream = nativeSql
+            .from(e)
+            .openStream()
+        assertFalse(config.dataSource.connection.closed)
+        stream.use { }
+        assertTrue(config.dataSource.connection.closed)
+    }
 }
diff --git a/integration-test-java/src/test/java/org/seasar/doma/it/criteria/NativeSqlSelectTest.java b/integration-test-java/src/test/java/org/seasar/doma/it/criteria/NativeSqlSelectTest.java
diff --git a/integration-test-kotlin/src/test/kotlin/org/seasar/doma/it/criteria/KNativeSqlSelectTest.kt b/integration-test-kotlin/src/test/kotlin/org/seasar/doma/it/criteria/KNativeSqlSelectTest.kt
