Stream support (#342)

adds `readStream`, `writeStream`, `closeStream` APIs for DBOS (and
`readStream` for DBOSClient) with associated DBOSExecutor and
SystemDatabase changes. Plus tests

fixes #215

Author: devhawk

transact/src/main/java/dev/dbos/transact/DBOS.java

@@ -35,6 +35,7 @@
 import java.time.ZoneId;
 import java.util.Collection;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -1107,4 +1108,49 @@ public static boolean inStep() {
   public static @Nullable SerializationStrategy serializationStrategy() {
     return DBOSContext.serializationStrategy();
   }
+
+  /**
+   * Write a value to a stream. Must be called from within a workflow or step.
+   *
+   * @param key The stream key / name within the workflow
+   * @param value A serializable value to write to the stream
+   */
+  public void writeStream(@NonNull String key, @NonNull Object value) {
+    writeStream(key, value, null);
+  }
+
+  /**
+   * Write a value to a stream with a specific serialization strategy. Must be called from within a
+   * workflow or step.
+   *
+   * @param key The stream key / name within the workflow
+   * @param value A serializable value to write to the stream
+   * @param serialization The serialization strategy to use (null for workflow default)
+   */
+  public void writeStream(
+      @NonNull String key, @NonNull Object value, @Nullable SerializationStrategy serialization) {
+    ensureLaunched("writeStream").writeStream(key, value, serialization);
+  }
+
+  /**
+   * Close a stream. Must be called from within a workflow, not a step.
+   *
+   * @param key The stream key / name within the workflow
+   */
+  public void closeStream(@NonNull String key) {
+    ensureLaunched("closeStream").closeStream(key);
+  }
+
+  /**
+   * Read values from a stream as an iterator. This function reads values from a stream identified
+   * by the workflow_id and key, returning an iterator that yields each value in order until the
+   * stream is closed or the workflow terminates.
+   *
+   * @param workflowId The workflow instance ID that owns the stream
+   * @param key The stream key / name within the workflow
+   * @return Iterator that yields each value in the stream
+   */
+  public @NonNull Iterator<Object> readStream(@NonNull String workflowId, @NonNull String key) {
+    return ensureLaunched("readStream").readStream(workflowId, key);
+  }
 }

transact/src/main/java/dev/dbos/transact/DBOSClient.java

@@ -1,6 +1,7 @@
 package dev.dbos.transact;
 
 import dev.dbos.transact.database.Result;
+import dev.dbos.transact.database.StreamIterator;
 import dev.dbos.transact.database.SystemDatabase;
 import dev.dbos.transact.execution.DBOSExecutor;
 import dev.dbos.transact.json.DBOSSerializer;
@@ -20,6 +21,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.time.ZoneId;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -642,6 +644,19 @@ public void send(
     return Optional.ofNullable(systemDatabase.getEvent(targetId, key, timeout, null));
   }
 
+  /**
+   * Read values from a stream as an iterator. This function reads values from a stream identified
+   * by the workflow_id and key, returning an iterator that yields each value in order until the
+   * stream is closed or the workflow terminates.
+   *
+   * @param workflowId The workflow instance ID that owns the stream
+   * @param key The stream key / name within the workflow
+   * @return Iterator that yields each value in the stream
+   */
+  public @NonNull Iterator<Object> readStream(@NonNull String workflowId, @NonNull String key) {
+    return new StreamIterator(workflowId, key, systemDatabase);
+  }
+
   /**
    * Create a handle for a workflow. This call does not ensure that the workflow exists; use the
    * returned handle's `getStatus()`.

transact/src/main/java/dev/dbos/transact/database/StreamIterator.java

@@ -0,0 +1,72 @@
+package dev.dbos.transact.database;
+
+import dev.dbos.transact.workflow.WorkflowState;
+import dev.dbos.transact.workflow.WorkflowStatus;
+
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+
+public class StreamIterator implements Iterator<Object> {
+  private final String workflowId;
+  private final String key;
+  private final SystemDatabase systemDatabase;
+  private int offset = 0;
+  private Object nextValue = null;
+  private boolean finished = false;
+
+  public StreamIterator(String workflowId, String key, SystemDatabase systemDatabase) {
+    this.workflowId = workflowId;
+    this.key = key;
+    this.systemDatabase = systemDatabase;
+    advance();
+  }
+
+  private void advance() {
+    while (!finished) {
+      try {
+        Object value = systemDatabase.readStream(workflowId, key, offset);
+        nextValue = value;
+        offset++;
+        return;
+      } catch (IllegalArgumentException e) {
+        WorkflowStatus status = systemDatabase.getWorkflowStatus(workflowId);
+        if (status == null || !isWorkflowActive(status.status())) {
+          finished = true;
+          nextValue = null;
+          return;
+        }
+        try {
+          Thread.sleep(1000);
+        } catch (InterruptedException ie) {
+          Thread.currentThread().interrupt();
+          finished = true;
+          nextValue = null;
+          return;
+        }
+      } catch (IllegalStateException e) {
+        finished = true;
+        nextValue = null;
+        return;
+      }
+    }
+  }
+
+  private boolean isWorkflowActive(WorkflowState state) {
+    return WorkflowState.PENDING == state || WorkflowState.ENQUEUED == state;
+  }
+
+  @Override
+  public boolean hasNext() {
+    return nextValue != null;
+  }
+
+  @Override
+  public Object next() {
+    if (nextValue == null) {
+      throw new NoSuchElementException();
+    }
+    Object result = nextValue;
+    advance();
+    return result;
+  }
+}

transact/src/main/java/dev/dbos/transact/database/StreamsDAO.java

@@ -0,0 +1,193 @@
+package dev.dbos.transact.database;
+
+import dev.dbos.transact.json.SerializationUtil;
+import dev.dbos.transact.workflow.internal.StepResult;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import javax.sql.DataSource;
+
+class StreamsDAO {
+
+  private final DataSource dataSource;
+  private final String schema;
+
+  StreamsDAO(DataSource dataSource, String schema) {
+    this.dataSource = dataSource;
+    this.schema = schema;
+  }
+
+  public void writeStreamFromStep(
+      String workflowId, int functionId, String key, Object value, String serializationFormat)
+      throws SQLException {
+    try (Connection conn = dataSource.getConnection()) {
+      insertStream(conn, workflowId, functionId, key, value, serializationFormat);
+    }
+  }
+
+  public void writeStreamFromWorkflow(
+      String workflowId, int functionId, String key, Object value, String serializationFormat)
+      throws SQLException {
+    String functionName =
+        STREAM_CLOSED_SENTINEL.equals(value) ? "DBOS.closeStream" : "DBOS.writeStream";
+    long startTime = System.currentTimeMillis();
+
+    try (Connection conn = dataSource.getConnection()) {
+      conn.setAutoCommit(false);
+
+      try {
+        StepResult recordedOutput =
+            StepsDAO.checkStepExecutionTxn(workflowId, functionId, functionName, conn, schema);
+
+        if (recordedOutput != null) {
+          logger.debug("Replaying writeStream, id: {}, key: {}", functionId, key);
+          conn.commit();
+          return;
+        } else {
+          logger.debug("Running writeStream, id: {}, key: {}", functionId, key);
+        }
+
+        insertStream(conn, workflowId, functionId, key, value, serializationFormat);
+
+        var output = new StepResult(workflowId, functionId, functionName, null, null, null, null);
+        StepsDAO.recordStepResultTxn(output, startTime, System.currentTimeMillis(), conn, schema);
+
+        conn.commit();
+
+      } catch (Exception e) {
+        try {
+          conn.rollback();
+        } catch (SQLException rollbackEx) {
+          e.addSuppressed(rollbackEx);
+        }
+        throw e;
+      }
+    }
+  }
+
+  private void insertStream(
+      Connection conn,
+      String workflowId,
+      int functionId,
+      String key,
+      Object value,
+      String serializationFormat)
+      throws SQLException {
+    var serialized = SerializationUtil.serializeValue(value, serializationFormat, null);
+    int offset = getNextOffsetTx(conn, workflowId, key);
+
+    String sql =
+        """
+        INSERT INTO "%s".streams (workflow_uuid, key, value, "offset", function_id, serialization)
+        VALUES (?, ?, ?, ?, ?, ?)
+        """
+            .formatted(schema);
+
+    try (var stmt = conn.prepareStatement(sql)) {
+      stmt.setString(1, workflowId);
+      stmt.setString(2, key);
+      stmt.setString(3, serialized.serializedValue());
+      stmt.setInt(4, offset);
+      stmt.setInt(5, functionId);
+      stmt.setString(6, serialized.serialization());
+      stmt.executeUpdate();
+    }
+  }
+
+  private int getNextOffsetTx(Connection conn, String workflowId, String key) throws SQLException {
+    String sql =
+        """
+        SELECT COALESCE(MAX("offset"), -1) + 1
+        FROM "%s".streams
+        WHERE workflow_uuid = ? AND key = ?
+        """
+            .formatted(schema);
+
+    try (var stmt = conn.prepareStatement(sql)) {
+      stmt.setString(1, workflowId);
+      stmt.setString(2, key);
+      try (var rs = stmt.executeQuery()) {
+        if (rs.next()) {
+          return rs.getInt(1);
+        }
+        return 0;
+      }
+    }
+  }
+
+  public void closeStream(String workflowId, int functionId, String key) throws SQLException {
+    writeStreamFromWorkflow(workflowId, functionId, key, STREAM_CLOSED_SENTINEL, "portable_json");
+  }
+
+  public Object readStream(String workflowId, String key, int offset) throws SQLException {
+    String sql =
+        """
+        SELECT value, serialization
+        FROM "%s".streams
+        WHERE workflow_uuid = ? AND key = ? AND "offset" = ?
+        """
+            .formatted(schema);
+
+    try (Connection conn = dataSource.getConnection();
+        var stmt = conn.prepareStatement(sql)) {
+      stmt.setString(1, workflowId);
+      stmt.setString(2, key);
+      stmt.setInt(3, offset);
+      try (var rs = stmt.executeQuery()) {
+        if (rs.next()) {
+          String value = rs.getString("value");
+          String serialization = rs.getString("serialization");
+          Object deserialized = SerializationUtil.deserializeValue(value, serialization, null);
+          if (STREAM_CLOSED_SENTINEL.equals(deserialized)) {
+            throw new IllegalStateException("Stream closed for key: " + key);
+          }
+          return deserialized;
+        }
+        throw new IllegalArgumentException(
+            "No value found for workflow=" + workflowId + ", key=" + key + ", offset=" + offset);
+      }
+    }
+  }
+
+  public Map<String, List<Object>> getAllStreamEntries(String workflowId) throws SQLException {
+    String sql =
+        """
+        SELECT key, value, serialization
+        FROM "%s".streams
+        WHERE workflow_uuid = ?
+        ORDER BY key, "offset"
+        """
+            .formatted(schema);
+
+    var streams = new LinkedHashMap<String, List<Object>>();
+    try (Connection conn = dataSource.getConnection();
+        var stmt = conn.prepareStatement(sql)) {
+      stmt.setString(1, workflowId);
+      try (var rs = stmt.executeQuery()) {
+        while (rs.next()) {
+          String key = rs.getString("key");
+          String value = rs.getString("value");
+          String serialization = rs.getString("serialization");
+
+          Object deserialized = SerializationUtil.deserializeValue(value, serialization, null);
+          if (STREAM_CLOSED_SENTINEL.equals(deserialized)) {
+            continue;
+          }
+
+          streams.computeIfAbsent(key, k -> new ArrayList<>()).add(deserialized);
+        }
+      }
+    }
+    return streams;
+  }
+
+  private static final org.slf4j.Logger logger =
+      org.slf4j.LoggerFactory.getLogger(StreamsDAO.class);
+
+  static final String STREAM_CLOSED_SENTINEL = "__DBOS_STREAM_CLOSED__";
+}