Register transactional step factories' checkpoint stores with DBOS

transact-jdbi-step-factory/src/main/java/dev/dbos/transact/jdbi/JdbiStepFactory.java

@@ -7,6 +7,7 @@
 import dev.dbos.transact.txstep.TxStepSchema;
 import dev.dbos.transact.workflow.internal.StepResult;
 
+import java.util.Collection;
 import java.util.Objects;
 import java.util.Optional;
 
@@ -24,6 +25,10 @@
  * Lambdas passed to {@link #inStep} or {@link #useStep} receive a {@link Handle} with a transaction
  * already open; they must not call {@code commit} or {@code close} themselves.
  *
+ * <p>Create the factory before calling {@link dev.dbos.transact.DBOS#launch()}: DBOS needs to know
+ * every factory when it launches, so the constructor throws {@link IllegalStateException} after
+ * launch.
+ *
  * <pre>{@code
  * JdbiStepFactory factory = new JdbiStepFactory(dbos, Jdbi.create(dataSource));
  *
@@ -232,6 +237,45 @@ protected void recordError(String workflowId, int stepId, Exception exception) {
     }
   }
 
+  @Override
+  protected void deleteCheckpoints(String workflowId, int fromStepId) {
+    inOwnTransaction(
+        h ->
+            h.createUpdate(TxStepSchema.deleteFromStepSql(schema))
+                .bind(0, workflowId)
+                .bind(1, fromStepId)
+                .execute());
+  }
+
+  @Override
+  protected void deleteCheckpoints(Collection<String> workflowIds) {
+    inOwnTransaction(
+        h ->
+            h.createUpdate(TxStepSchema.deleteWorkflowsSql(schema))
+                .bindArray(0, String.class, workflowIds)
+                .execute());
+  }
+
+  // Jdbi's useTransaction does not commit a connection that arrives with autocommit off: it takes
+  // that for a transaction already in progress and leaves the commit to someone else. begin() and
+  // commit() on the handle commit either way.
+  private void inOwnTransaction(HandleConsumer<RuntimeException> work) {
+    try (var handle = jdbi.open()) {
+      handle.begin();
+      try {
+        work.useHandle(handle);
+        handle.commit();
+      } catch (RuntimeException e) {
+        try {
+          handle.rollback();
+        } catch (RuntimeException rollbackFailure) {
+          e.addSuppressed(rollbackFailure);
+        }
+        throw e;
+      }
+    }
+  }
+
   private void recordResult(
       Handle handle,
       String workflowId,

transact-jdbi-step-factory/src/test/java/dev/dbos/transact/DBOSTestAccess.java

@@ -0,0 +1,11 @@
+package dev.dbos.transact;
+
+import dev.dbos.transact.execution.DBOSExecutor;
+
+// Helper class to retrieve DBOS internals via package private methods
+public class DBOSTestAccess {
+
+  public static DBOSExecutor getDbosExecutor(DBOS dbos) {
+    return dbos.getDbosExecutor();
+  }
+}

transact-jdbi-step-factory/src/test/java/dev/dbos/transact/jdbi/JdbiStepFactoryNoAutoCommitTest.java

@@ -0,0 +1,91 @@
+package dev.dbos.transact.jdbi;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import dev.dbos.transact.DBOS;
+import dev.dbos.transact.DBOSTestAccess;
+import dev.dbos.transact.utils.PgContainer;
+
+import java.sql.SQLException;
+import java.util.List;
+
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+import org.jdbi.v3.core.Jdbi;
+import org.junit.jupiter.api.AutoClose;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * JdbiStepFactory on a pool configured without autocommit. A delete that did not commit explicitly
+ * would be rolled back when its connection went back to the pool.
+ */
+public class JdbiStepFactoryNoAutoCommitTest {
+  private static final String SCHEMA = "jdbi_no_autocommit";
+
+  @AutoClose final PgContainer pgContainer = new PgContainer();
+
+  @AutoClose DBOS dbos;
+  @AutoClose HikariDataSource dataSource;
+  @AutoClose HikariDataSource noAutoCommitDataSource;
+
+  @BeforeEach
+  void beforeEach() {
+    dataSource = pgContainer.dataSource();
+    var cfg = new HikariConfig();
+    cfg.setJdbcUrl(pgContainer.jdbcUrl());
+    cfg.setUsername(pgContainer.username());
+    cfg.setPassword(pgContainer.password());
+    cfg.setAutoCommit(false);
+    noAutoCommitDataSource = new HikariDataSource(cfg);
+
+    dbos = new DBOS(pgContainer.dbosConfig());
+    new JdbiStepFactory(dbos, Jdbi.create(noAutoCommitDataSource), SCHEMA);
+    dbos.launch();
+  }
+
+  private void insertCheckpoints(String workflowId, int... stepIds) throws SQLException {
+    try (var conn = dataSource.getConnection();
+        var stmt =
+            conn.prepareStatement(
+                "INSERT INTO %s.tx_step_outputs (workflow_id, step_id) VALUES (?, ?)"
+                    .formatted(SCHEMA))) {
+      for (var stepId : stepIds) {
+        stmt.setString(1, workflowId);
+        stmt.setInt(2, stepId);
+        stmt.executeUpdate();
+      }
+    }
+  }
+
+  private int checkpoints(String workflowId) throws SQLException {
+    try (var conn = dataSource.getConnection();
+        var stmt =
+            conn.prepareStatement(
+                "SELECT COUNT(*) FROM %s.tx_step_outputs WHERE workflow_id = ?"
+                    .formatted(SCHEMA))) {
+      stmt.setString(1, workflowId);
+      try (var rs = stmt.executeQuery()) {
+        rs.next();
+        return rs.getInt(1);
+      }
+    }
+  }
+
+  @Test
+  public void deletesCommit() throws Exception {
+    insertCheckpoints("wf-one", 0, 1, 2);
+    insertCheckpoints("wf-a", 0);
+    insertCheckpoints("wf-b", 0);
+    insertCheckpoints("wf-keep", 0);
+    var executor = DBOSTestAccess.getDbosExecutor(dbos);
+
+    executor.deleteStepCheckpoints("wf-one", 1);
+    assertEquals(1, checkpoints("wf-one"));
+
+    executor.deleteStepCheckpoints(List.of("wf-a", "wf-b"));
+    assertEquals(0, checkpoints("wf-a"));
+    assertEquals(0, checkpoints("wf-b"));
+    assertEquals(1, checkpoints("wf-keep"));
+  }
+}

transact-jdbi-step-factory/src/test/java/dev/dbos/transact/jdbi/JdbiStepFactoryTest.java

@@ -7,6 +7,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import dev.dbos.transact.DBOS;
+import dev.dbos.transact.DBOSTestAccess;
 import dev.dbos.transact.config.DBOSConfig;
 import dev.dbos.transact.context.WorkflowOptions;
 import dev.dbos.transact.database.SystemDatabase;
@@ -17,6 +18,7 @@
 import dev.dbos.transact.workflow.WorkflowHandle;
 
 import java.sql.SQLException;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -228,6 +230,30 @@ private int getGreetCount(String user) throws SQLException {
     }
   }
 
+  @Test
+  public void testDeleteCheckpointsThroughDbos() throws Exception {
+    var wfid = "wf-delete";
+    try (var _o = new WorkflowOptions(wfid).setContext()) {
+      proxy.insertWorkflow("deleteUser");
+    }
+    assertEquals(1, DBUtils.getTxStepRows(dataSource, wfid).size());
+
+    // DBOS reaches the factory's checkpoints through the store it registered, and the delete
+    // commits through the library's own transaction.
+    DBOSTestAccess.getDbosExecutor(dbos).deleteStepCheckpoints(wfid, 0);
+    assertEquals(0, DBUtils.getTxStepRows(dataSource, wfid).size());
+
+    // The batch form, which binds the workflow IDs as an array.
+    for (var id : List.of("wf-batch-1", "wf-batch-2")) {
+      try (var _o = new WorkflowOptions(id).setContext()) {
+        proxy.insertWorkflow("deleteUser");
+      }
+    }
+    DBOSTestAccess.getDbosExecutor(dbos).deleteStepCheckpoints(List.of("wf-batch-1", "wf-batch-2"));
+    assertEquals(0, DBUtils.getTxStepRows(dataSource, "wf-batch-1").size());
+    assertEquals(0, DBUtils.getTxStepRows(dataSource, "wf-batch-2").size());
+  }
+
   @Test
   public void testInsert() throws Exception {
     var wfid = "wf1";

transact-jooq-step-factory/src/main/java/dev/dbos/transact/jooq/JooqStepFactory.java

@@ -9,6 +9,7 @@
 import dev.dbos.transact.txstep.TxStepSchema;
 import dev.dbos.transact.workflow.internal.StepResult;
 
+import java.util.Collection;
 import java.util.Objects;
 import java.util.Optional;
 
@@ -27,6 +28,10 @@
  * org.jooq.Configuration} with a transaction already open; they must not commit or close the
  * underlying connection themselves.
  *
+ * <p>Create the factory before calling {@link dev.dbos.transact.DBOS#launch()}: DBOS needs to know
+ * every factory when it launches, so the constructor throws {@link IllegalStateException} after
+ * launch.
+ *
  * <pre>{@code
  * JooqStepFactory factory = new JooqStepFactory(dbos, dslContext);
  *
@@ -208,6 +213,21 @@ protected void recordError(String workflowId, int stepId, Exception exception) {
     }
   }
 
+  @Override
+  protected void deleteCheckpoints(String workflowId, int fromStepId) {
+    dsl.transaction(
+        trx -> trx.dsl().execute(TxStepSchema.deleteFromStepSql(schema), workflowId, fromStepId));
+  }
+
+  @Override
+  protected void deleteCheckpoints(Collection<String> workflowIds) {
+    // Bound on the JDBC connection: jOOQ can only bind an array when the dialect is set to one
+    // that has arrays, and the DSLContext may have been built without one.
+    dsl.transaction(
+        trx ->
+            trx.dsl().connection(conn -> TxStepSchema.deleteWorkflows(conn, schema, workflowIds)));
+  }
+
   private void recordResult(
       DSLContext ctx,
       String workflowId,

transact-jooq-step-factory/src/test/java/dev/dbos/transact/DBOSTestAccess.java

@@ -0,0 +1,11 @@
+package dev.dbos.transact;
+
+import dev.dbos.transact.execution.DBOSExecutor;
+
+// Helper class to retrieve DBOS internals via package private methods
+public class DBOSTestAccess {
+
+  public static DBOSExecutor getDbosExecutor(DBOS dbos) {
+    return dbos.getDbosExecutor();
+  }
+}

transact-jooq-step-factory/src/test/java/dev/dbos/transact/jooq/JooqStepFactoryNoAutoCommitTest.java

@@ -0,0 +1,93 @@
+package dev.dbos.transact.jooq;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import dev.dbos.transact.DBOS;
+import dev.dbos.transact.DBOSTestAccess;
+import dev.dbos.transact.utils.PgContainer;
+
+import java.sql.SQLException;
+import java.util.List;
+
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+import org.jooq.SQLDialect;
+import org.jooq.impl.DSL;
+import org.junit.jupiter.api.AutoClose;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * JooqStepFactory on a pool configured without autocommit, and with jOOQ's default dialect, which
+ * cannot bind arrays itself. A delete that did not commit explicitly would be rolled back when its
+ * connection went back to the pool.
+ */
+public class JooqStepFactoryNoAutoCommitTest {
+  private static final String SCHEMA = "jooq_no_autocommit";
+
+  @AutoClose final PgContainer pgContainer = new PgContainer();
+
+  @AutoClose DBOS dbos;
+  @AutoClose HikariDataSource dataSource;
+  @AutoClose HikariDataSource noAutoCommitDataSource;
+
+  @BeforeEach
+  void beforeEach() {
+    dataSource = pgContainer.dataSource();
+    var cfg = new HikariConfig();
+    cfg.setJdbcUrl(pgContainer.jdbcUrl());
+    cfg.setUsername(pgContainer.username());
+    cfg.setPassword(pgContainer.password());
+    cfg.setAutoCommit(false);
+    noAutoCommitDataSource = new HikariDataSource(cfg);
+
+    dbos = new DBOS(pgContainer.dbosConfig());
+    new JooqStepFactory(dbos, DSL.using(noAutoCommitDataSource, SQLDialect.DEFAULT), SCHEMA);
+    dbos.launch();
+  }
+
+  private void insertCheckpoints(String workflowId, int... stepIds) throws SQLException {
+    try (var conn = dataSource.getConnection();
+        var stmt =
+            conn.prepareStatement(
+                "INSERT INTO %s.tx_step_outputs (workflow_id, step_id) VALUES (?, ?)"
+                    .formatted(SCHEMA))) {
+      for (var stepId : stepIds) {
+        stmt.setString(1, workflowId);
+        stmt.setInt(2, stepId);
+        stmt.executeUpdate();
+      }
+    }
+  }
+
+  private int checkpoints(String workflowId) throws SQLException {
+    try (var conn = dataSource.getConnection();
+        var stmt =
+            conn.prepareStatement(
+                "SELECT COUNT(*) FROM %s.tx_step_outputs WHERE workflow_id = ?"
+                    .formatted(SCHEMA))) {
+      stmt.setString(1, workflowId);
+      try (var rs = stmt.executeQuery()) {
+        rs.next();
+        return rs.getInt(1);
+      }
+    }
+  }
+
+  @Test
+  public void deletesCommit() throws Exception {
+    insertCheckpoints("wf-one", 0, 1, 2);
+    insertCheckpoints("wf-a", 0);
+    insertCheckpoints("wf-b", 0);
+    insertCheckpoints("wf-keep", 0);
+    var executor = DBOSTestAccess.getDbosExecutor(dbos);
+
+    executor.deleteStepCheckpoints("wf-one", 1);
+    assertEquals(1, checkpoints("wf-one"));
+
+    executor.deleteStepCheckpoints(List.of("wf-a", "wf-b"));
+    assertEquals(0, checkpoints("wf-a"));
+    assertEquals(0, checkpoints("wf-b"));
+    assertEquals(1, checkpoints("wf-keep"));
+  }
+}