Handle DBOS operations as workflow steps (#120)

Call DBOS operations inside `RunAsStep` if within a workflow.

One thing this PR had to fix is the forked workflow ID generation, which
was at a higher level and not idempotent. The systemDB forkWorkflow now
returns the forked workflow ID.

Author: maxdml

dbos/dbos.go

@@ -318,6 +318,10 @@ func NewDBOSContext(ctx context.Context, inputConfig Config) (DBOSContext, error
 	// Register types we serialize with gob
 	var t time.Time
 	gob.Register(t)
+	var ws []WorkflowStatus
+	gob.Register(ws)
+	var si []StepInfo
+	gob.Register(si)
 
 	// Initialize global variables from processed config (already handles env vars and defaults)
 	initExecutor.applicationVersion = config.ApplicationVersion

dbos/queues_test.go

@@ -926,23 +926,17 @@ func TestQueueTimeouts(t *testing.T) {
 	}
 	RegisterWorkflow(dbosCtx, queuedWaitForCancelWorkflow)
 
-	enqueuedWorkflowEnqueuesATimeoutWorkflow := func(ctx DBOSContext, _ string) (string, error) {
+	enqueuedWorkflowEnqueuesATimeoutWorkflow := func(ctx DBOSContext, childWorkflowID string) (string, error) {
 		// This workflow will enqueue a workflow that waits indefinitely until it is cancelled
-		handle, err := RunWorkflow(ctx, queuedWaitForCancelWorkflow, "enqueued-wait-for-cancel", WithQueue(timeoutQueue.Name))
+		handle, err := RunWorkflow(ctx, queuedWaitForCancelWorkflow, "enqueued-wait-for-cancel", WithQueue(timeoutQueue.Name), WithWorkflowID(childWorkflowID))
 		require.NoError(t, err, "failed to start enqueued wait for cancel workflow")
 		// Workflow should get AwaitedWorkflowCancelled DBOSError
 		_, err = handle.GetResult()
 		require.Error(t, err, "expected error when waiting for enqueued workflow to complete, but got none")
 		var dbosErr *DBOSError
 		require.ErrorAs(t, err, &dbosErr, "expected error to be of type *DBOSError, got %T", err)
 		assert.Equal(t, AwaitedWorkflowCancelled, dbosErr.Code, "expected error code to be AwaitedWorkflowCancelled")
-
-		// enqueud workflow should have been cancelled
-		status, err := handle.GetStatus()
-		require.NoError(t, err, "failed to get status of enqueued workflow")
-		assert.Equal(t, WorkflowStatusCancelled, status.Status, "expected enqueued workflow status to be WorkflowStatusCancelled")
-
-		return "should-never-see-this", nil
+		return "", nil
 	}
 	RegisterWorkflow(dbosCtx, enqueuedWorkflowEnqueuesATimeoutWorkflow)
 
@@ -956,19 +950,20 @@ func TestQueueTimeouts(t *testing.T) {
 	}
 
 	enqueuedWorkflowEnqueuesADetachedWorkflow := func(ctx DBOSContext, timeout time.Duration) (string, error) {
+		myId, err := GetWorkflowID(ctx)
+		if err != nil {
+			return "", fmt.Errorf("failed to get workflow ID: %v", err)
+		}
+		childID := fmt.Sprintf("%s-child", myId)
 		// This workflow will enqueue a workflow that is not cancelable
 		childCtx := WithoutCancel(ctx)
-		handle, err := RunWorkflow(childCtx, detachedWorkflow, timeout*2, WithQueue(timeoutQueue.Name))
+		handle, err := RunWorkflow(childCtx, detachedWorkflow, timeout*2, WithQueue(timeoutQueue.Name), WithWorkflowID(childID))
 		require.NoError(t, err, "failed to start enqueued detached workflow")
 
 		// Wait for the enqueued workflow to complete
 		result, err := handle.GetResult()
 		require.NoError(t, err, "failed to get result from enqueued detached workflow")
 		assert.Equal(t, "detached-workflow-completed", result, "expected result to be 'detached-workflow-completed'")
-		// Check the workflow status: should be success
-		status, err := handle.GetStatus()
-		require.NoError(t, err, "failed to get enqueued detached workflow status")
-		assert.Equal(t, WorkflowStatusSuccess, status.Status, "expected enqueued detached workflow status to be WorkflowStatusSuccess")
 		return result, nil
 	}
 
@@ -1022,7 +1017,8 @@ func TestQueueTimeouts(t *testing.T) {
 		cancelCtx, cancelFunc := WithTimeout(dbosCtx, 1*time.Millisecond)
 		defer cancelFunc() // Ensure we clean up the context
 
-		handle, err := RunWorkflow(cancelCtx, enqueuedWorkflowEnqueuesATimeoutWorkflow, "enqueue-timeout-workflow", WithQueue(timeoutQueue.Name))
+		childWorkflowID := uuid.NewString()
+		handle, err := RunWorkflow(cancelCtx, enqueuedWorkflowEnqueuesATimeoutWorkflow, childWorkflowID, WithQueue(timeoutQueue.Name))
 		require.NoError(t, err, "failed to start enqueued workflow")
 
 		// Wait for the workflow to complete and get the result
@@ -1042,6 +1038,18 @@ func TestQueueTimeouts(t *testing.T) {
 		require.NoError(t, err, "failed to get workflow status")
 		assert.Equal(t, WorkflowStatusCancelled, status.Status, "expected workflow status to be WorkflowStatusCancelled")
 
+		// Wait for the child workflow status to become cancelled
+		require.Eventually(t, func() bool {
+			childHandle, err := RetrieveWorkflow[string](dbosCtx, childWorkflowID)
+			require.NoError(t, err, "failed to retrieve child workflow")
+
+			status, err := childHandle.GetStatus()
+			if err != nil {
+				return false
+			}
+			return status.Status == WorkflowStatusCancelled
+		}, 5*time.Second, 100*time.Millisecond, "expected enqueued workflow status to be WorkflowStatusCancelled")
+
 		require.True(t, queueEntriesAreCleanedUp(dbosCtx), "expected queue entries to be cleaned up after workflow cancellation, but they are not")
 	})
 
@@ -1071,6 +1079,19 @@ func TestQueueTimeouts(t *testing.T) {
 		require.NoError(t, err, "failed to get enqueued detached workflow status")
 		assert.Equal(t, WorkflowStatusCancelled, status.Status, "expected enqueued detached workflow status to be WorkflowStatusCancelled")
 
+		// Check the child's status: should be success because it is detached
+		require.Eventually(t, func() bool {
+			childID := fmt.Sprintf("%s-child", handle.GetWorkflowID())
+			childHandle, err := RetrieveWorkflow[string](dbosCtx, childID)
+			require.NoError(t, err, "failed to retrieve detached workflow")
+
+			status, err := childHandle.GetStatus()
+			if err != nil {
+				return false
+			}
+			return status.Status == WorkflowStatusSuccess
+		}, 5*time.Second, 100*time.Millisecond, "expected detached workflow status to be WorkflowStatusSuccess")
+
 		require.True(t, queueEntriesAreCleanedUp(dbosCtx), "expected queue entries to be cleaned up after workflow cancellation, but they are not")
 	})
 

dbos/serialization_test.go

@@ -96,9 +96,12 @@ func TestWorkflowEncoding(t *testing.T) {
 		assert.Equal(t, "workflow error: step error", err.Error())
 
 		// Test results from ListWorkflows
-		workflows, err := ListWorkflows(executor, WithWorkflowIDs(
-			[]string{directHandle.GetWorkflowID()},
-		))
+		workflows, err := ListWorkflows(
+			executor,
+			WithWorkflowIDs([]string{directHandle.GetWorkflowID()}),
+			WithLoadInput(true),
+			WithLoadOutput(true),
+		)
 		require.NoError(t, err)
 		require.Len(t, workflows, 1)
 		workflow := workflows[0]
@@ -155,9 +158,12 @@ func TestWorkflowEncoding(t *testing.T) {
 		assert.Equal(t, "processed by encodingStepStruct", retrievedResult.B)
 
 		// Test results from ListWorkflows
-		workflows, err := ListWorkflows(executor, WithWorkflowIDs(
-			[]string{directHandle.GetWorkflowID()},
-		))
+		workflows, err := ListWorkflows(executor,
+			WithWorkflowIDs([]string{directHandle.GetWorkflowID()}),
+			WithLoadInput(true),
+			WithLoadOutput(true),
+		)
+		require.Len(t, workflows, 1)
 		require.NoError(t, err)
 		workflow := workflows[0]
 		require.NotNil(t, workflow.Input)

dbos/system_database.go

@@ -12,6 +12,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -35,7 +36,7 @@ type systemDatabase interface {
 	cancelWorkflow(ctx context.Context, workflowID string) error
 	cancelAllBefore(ctx context.Context, cutoffTime time.Time) error
 	resumeWorkflow(ctx context.Context, workflowID string) error
-	forkWorkflow(ctx context.Context, input forkWorkflowDBInput) error
+	forkWorkflow(ctx context.Context, input forkWorkflowDBInput) (string, error)
 
 	// Child workflows
 	recordChildWorkflow(ctx context.Context, input recordChildWorkflowDBInput) error
@@ -657,7 +658,7 @@ func (s *sysDB) listWorkflows(ctx context.Context, input listWorkflowsDBInput) (
 		var executorID *string
 
 		// Build scan arguments dynamically based on loaded columns
-		scanArgs := []interface{}{
+		scanArgs := []any{
 			&wf.ID, &wf.Status, &wf.Name, &wf.AuthenticatedUser, &wf.AssumedRole,
 			&wf.AuthenticatedRoles, &executorID, &createdAtMs,
 			&updatedAtMs, &applicationVersion, &wf.ApplicationID,
@@ -849,7 +850,6 @@ func (s *sysDB) cancelAllBefore(ctx context.Context, cutoffTime time.Time) error
 			// If desired we could funnel the errors back the caller (conductor, admin server)
 		}
 	}
-
 	return nil
 }
 
@@ -977,15 +977,21 @@ type forkWorkflowDBInput struct {
 	applicationVersion string
 }
 
-func (s *sysDB) forkWorkflow(ctx context.Context, input forkWorkflowDBInput) error {
+func (s *sysDB) forkWorkflow(ctx context.Context, input forkWorkflowDBInput) (string, error) {
+	// Generate new workflow ID if not provided
+	forkedWorkflowID := input.forkedWorkflowID
+	if forkedWorkflowID == "" {
+		forkedWorkflowID = uuid.New().String()
+	}
+
 	// Validate startStep
 	if input.startStep < 0 {
-		return fmt.Errorf("startStep must be >= 0, got %d", input.startStep)
+		return "", fmt.Errorf("startStep must be >= 0, got %d", input.startStep)
 	}
 
 	tx, err := s.pool.Begin(ctx)
 	if err != nil {
-		return fmt.Errorf("failed to begin transaction: %w", err)
+		return "", fmt.Errorf("failed to begin transaction: %w", err)
 	}
 	defer tx.Rollback(ctx)
 
@@ -997,10 +1003,10 @@ func (s *sysDB) forkWorkflow(ctx context.Context, input forkWorkflowDBInput) err
 	}
 	wfs, err := s.listWorkflows(ctx, listInput)
 	if err != nil {
-		return fmt.Errorf("failed to list workflows: %w", err)
+		return "", fmt.Errorf("failed to list workflows: %w", err)
 	}
 	if len(wfs) == 0 {
-		return newNonExistentWorkflowError(input.originalWorkflowID)
+		return "", newNonExistentWorkflowError(input.originalWorkflowID)
 	}
 
 	originalWorkflow := wfs[0]
@@ -1030,11 +1036,11 @@ func (s *sysDB) forkWorkflow(ctx context.Context, input forkWorkflowDBInput) err
 
 	inputString, err := serialize(originalWorkflow.Input)
 	if err != nil {
-		return fmt.Errorf("failed to serialize input: %w", err)
+		return "", fmt.Errorf("failed to serialize input: %w", err)
 	}
 
 	_, err = tx.Exec(ctx, insertQuery,
-		input.forkedWorkflowID,
+		forkedWorkflowID,
 		WorkflowStatusEnqueued,
 		originalWorkflow.Name,
 		originalWorkflow.AuthenticatedUser,
@@ -1049,7 +1055,7 @@ func (s *sysDB) forkWorkflow(ctx context.Context, input forkWorkflowDBInput) err
 		0)
 
 	if err != nil {
-		return fmt.Errorf("failed to insert forked workflow status: %w", err)
+		return "", fmt.Errorf("failed to insert forked workflow status: %w", err)
 	}
 
 	// If startStep > 0, copy the original workflow's outputs into the forked workflow
@@ -1060,17 +1066,17 @@ func (s *sysDB) forkWorkflow(ctx context.Context, input forkWorkflowDBInput) err
 			FROM dbos.operation_outputs
 			WHERE workflow_uuid = $2 AND function_id < $3`
 
-		_, err = tx.Exec(ctx, copyOutputsQuery, input.forkedWorkflowID, input.originalWorkflowID, input.startStep)
+		_, err = tx.Exec(ctx, copyOutputsQuery, forkedWorkflowID, input.originalWorkflowID, input.startStep)
 		if err != nil {
-			return fmt.Errorf("failed to copy operation outputs: %w", err)
+			return "", fmt.Errorf("failed to copy operation outputs: %w", err)
 		}
 	}
 
 	if err := tx.Commit(ctx); err != nil {
-		return fmt.Errorf("failed to commit transaction: %w", err)
+		return "", fmt.Errorf("failed to commit transaction: %w", err)
 	}
 
-	return nil
+	return forkedWorkflowID, nil
 }
 
 func (s *sysDB) awaitWorkflowResult(ctx context.Context, workflowID string) (any, error) {