Merge branch 'main' into feat/register-workflowhandle

Author: maxdml

cmd/dbos/migrate.go

@@ -29,7 +29,7 @@ func init() {
 
 func runMigrate(cmd *cobra.Command, args []string) error {
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}

cmd/dbos/reset.go

@@ -34,7 +34,7 @@ func runReset(cmd *cobra.Command, args []string) error {
 	}
 
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}

cmd/dbos/utils.go

@@ -9,7 +9,6 @@ import (
 	"os"
 
 	"github.com/dbos-inc/dbos-transact-golang/dbos"
-	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
 
@@ -43,7 +42,7 @@ func maskPassword(dbURL string) string {
 }
 
 // getDBURL resolves the database URL from flag, config, or environment variable
-func getDBURL(_ *cobra.Command) (string, error) {
+func getDBURL() (string, error) {
 	var resolvedURL string
 	var source string
 

cmd/dbos/workflow.go

@@ -85,7 +85,7 @@ func init() {
 
 func runWorkflowList(cmd *cobra.Command, args []string) error {
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}
@@ -192,7 +192,7 @@ func runWorkflowGet(cmd *cobra.Command, args []string) error {
 	workflowID := args[0]
 
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}
@@ -227,7 +227,7 @@ func runWorkflowSteps(cmd *cobra.Command, args []string) error {
 	workflowID := args[0]
 
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}
@@ -259,7 +259,7 @@ func runWorkflowCancel(cmd *cobra.Command, args []string) error {
 	workflowID := args[0]
 
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}
@@ -285,7 +285,7 @@ func runWorkflowResume(cmd *cobra.Command, args []string) error {
 	workflowID := args[0]
 
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}
@@ -318,7 +318,7 @@ func runWorkflowFork(cmd *cobra.Command, args []string) error {
 	workflowID := args[0]
 
 	// Get database URL
-	dbURL, err := getDBURL(cmd)
+	dbURL, err := getDBURL()
 	if err != nil {
 		return err
 	}

dbos/client.go

@@ -2,7 +2,6 @@ package dbos
 
 import (
 	"context"
-	"encoding/gob"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -32,6 +31,7 @@ type Client interface {
 	CancelWorkflow(workflowID string) error
 	ResumeWorkflow(workflowID string) (WorkflowHandle[any], error)
 	ForkWorkflow(input ForkWorkflowInput) (WorkflowHandle[any], error)
+	GetWorkflowSteps(workflowID string) ([]StepInfo, error)
 	Shutdown(timeout time.Duration) // Simply close the system DB connection pool
 }
 
@@ -241,10 +241,16 @@ func Enqueue[P any, R any](c Client, queueName, workflowName string, input P, op
 	}
 
 	// Register the input and outputs for gob encoding
+	var logger *slog.Logger
+	if cl, ok := c.(*client); ok {
+		if ctx, ok := cl.dbosCtx.(*dbosContext); ok {
+			logger = ctx.logger
+		}
+	}
 	var typedInput P
-	gob.Register(typedInput)
+	safeGobRegister(typedInput, logger)
 	var typedOutput R
-	gob.Register(typedOutput)
+	safeGobRegister(typedOutput, logger)
 
 	// Call the interface method with the same signature
 	handle, err := c.Enqueue(queueName, workflowName, input, opts...)
@@ -290,6 +296,11 @@ func (c *client) ForkWorkflow(input ForkWorkflowInput) (WorkflowHandle[any], err
 	return c.dbosCtx.ForkWorkflow(c.dbosCtx, input)
 }
 
+// GetWorkflowSteps retrieves the execution steps of a workflow.
+func (c *client) GetWorkflowSteps(workflowID string) ([]StepInfo, error) {
+	return c.dbosCtx.GetWorkflowSteps(c.dbosCtx, workflowID)
+}
+
 // Shutdown gracefully shuts down the client and closes the system database connection.
 func (c *client) Shutdown(timeout time.Duration) {
 	// Get the concrete dbosContext to access internal fields

dbos/client_test.go

@@ -928,3 +928,69 @@ func TestListWorkflows(t *testing.T) {
 	// Verify all queue entries are cleaned up
 	require.True(t, queueEntriesAreCleanedUp(serverCtx), "expected queue entries to be cleaned up after list workflows tests")
 }
+
+func TestGetWorkflowSteps(t *testing.T) {
+	// Setup server context
+	serverCtx := setupDBOS(t, true, true)
+
+	// Create queue for communication
+	queue := NewWorkflowQueue(serverCtx, "get-workflow-steps-queue")
+
+	// Workflow with one step
+	stepFunction := func(ctx context.Context) (string, error) {
+		return "abc", nil
+	}
+
+	testWorkflow := func(ctx DBOSContext, input string) (string, error) {
+		result, err := RunAsStep(ctx, stepFunction, WithStepName("TestStep"))
+		if err != nil {
+			return "", err
+		}
+		return result, nil
+	}
+	RegisterWorkflow(serverCtx, testWorkflow, WithWorkflowName("TestWorkflow"))
+
+	// Launch server
+	err := Launch(serverCtx)
+	require.NoError(t, err)
+
+	// Setup client
+	databaseURL := getDatabaseURL()
+	config := ClientConfig{
+		DatabaseURL: databaseURL,
+	}
+	client, err := NewClient(context.Background(), config)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		if client != nil {
+			client.Shutdown(30 * time.Second)
+		}
+	})
+
+	// Enqueue and run the workflow
+	workflowID := "test-get-workflow-steps"
+	handle, err := Enqueue[string, string](client, queue.Name, "TestWorkflow", "test-input", WithEnqueueWorkflowID(workflowID))
+	require.NoError(t, err)
+
+	// Wait for workflow to complete
+	result, err := handle.GetResult()
+	require.NoError(t, err)
+	assert.Equal(t, "abc", result)
+
+	// Test GetWorkflowSteps with loadOutput = true
+	stepsWithOutput, err := client.GetWorkflowSteps(workflowID)
+	require.NoError(t, err)
+	require.Len(t, stepsWithOutput, 1, "expected exactly 1 step")
+
+	step := stepsWithOutput[0]
+	assert.Equal(t, 0, step.StepID, "expected step ID to be 0")
+	assert.Equal(t, "TestStep", step.StepName, "expected step name to be set")
+	assert.Nil(t, step.Error, "expected no error in step")
+	assert.Equal(t, "", step.ChildWorkflowID, "expected no child workflow ID")
+
+	// Verify the output wasn't loaded
+	require.Nil(t, step.Output, "expected output not to be loaded")
+
+	// Verify all queue entries are cleaned up
+	require.True(t, queueEntriesAreCleanedUp(serverCtx), "expected queue entries to be cleaned up after get workflow steps test")
+}

dbos/dbos.go

@@ -3,7 +3,6 @@ package dbos
 import (
 	"context"
 	"crypto/sha256"
-	"encoding/gob"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -121,12 +120,13 @@ type DBOSContext interface {
 	GetStepID() (int, error)                                                                                    // Get the current step ID (only available within workflows)
 
 	// Workflow management
-	RetrieveWorkflow(_ DBOSContext, workflowID string) (WorkflowHandle[any], error)     // Get a handle to an existing workflow
-	CancelWorkflow(_ DBOSContext, workflowID string) error                              // Cancel a workflow by setting its status to CANCELLED
-	ResumeWorkflow(_ DBOSContext, workflowID string) (WorkflowHandle[any], error)       // Resume a cancelled workflow
-	ForkWorkflow(_ DBOSContext, input ForkWorkflowInput) (WorkflowHandle[any], error)   // Fork a workflow from a specific step
-	ListWorkflows(_ DBOSContext, opts ...ListWorkflowsOption) ([]WorkflowStatus, error) // List workflows based on filtering criteria
-	GetWorkflowSteps(_ DBOSContext, workflowID string) ([]StepInfo, error)              // Get the execution steps of a workflow
+	RetrieveWorkflow(_ DBOSContext, workflowID string) (WorkflowHandle[any], error)                                // Get a handle to an existing workflow
+	CancelWorkflow(_ DBOSContext, workflowID string) error                                                         // Cancel a workflow by setting its status to CANCELLED
+	ResumeWorkflow(_ DBOSContext, workflowID string) (WorkflowHandle[any], error)                                  // Resume a cancelled workflow
+	ForkWorkflow(_ DBOSContext, input ForkWorkflowInput) (WorkflowHandle[any], error)                              // Fork a workflow from a specific step
+	ListWorkflows(_ DBOSContext, opts ...ListWorkflowsOption) ([]WorkflowStatus, error)                            // List workflows based on filtering criteria
+	GetWorkflowSteps(_ DBOSContext, workflowID string) ([]StepInfo, error)                                         // Get the execution steps of a workflow
+	ListRegisteredWorkflows(_ DBOSContext, opts ...ListRegisteredWorkflowsOption) ([]WorkflowRegistryEntry, error) // List registered workflows with filtering options
 
 	// Accessors
 	GetApplicationVersion() string // Get the application version for this context
@@ -159,7 +159,7 @@ type dbosContext struct {
 	workflowsWg *sync.WaitGroup
 
 	// Workflow registry - read-mostly sync.Map since registration happens only before launch
-	workflowRegistry        *sync.Map // map[string]workflowRegistryEntry
+	workflowRegistry        *sync.Map // map[string]WorkflowRegistryEntry
 	workflowCustomNametoFQN *sync.Map // Maps fully qualified workflow names to custom names. Usefor when client enqueues a workflow by name because registry is indexed by FQN.
 
 	// Workflow scheduler
@@ -194,7 +194,8 @@ func WithValue(ctx DBOSContext, key, val any) DBOSContext {
 	}
 	// Will do nothing if the concrete type is not dbosContext
 	if dbosCtx, ok := ctx.(*dbosContext); ok {
-		return &dbosContext{
+		launched := dbosCtx.launched.Load()
+		childCtx := &dbosContext{
 			ctx:                     context.WithValue(dbosCtx.ctx, key, val), // Spawn a new child context with the value set
 			logger:                  dbosCtx.logger,
 			systemDB:                dbosCtx.systemDB,
@@ -205,6 +206,8 @@ func WithValue(ctx DBOSContext, key, val any) DBOSContext {
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
 		}
+		childCtx.launched.Store(launched)
+		return childCtx
 	}
 	return nil
 }
@@ -217,7 +220,10 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 		return nil
 	}
 	if dbosCtx, ok := ctx.(*dbosContext); ok {
-		return &dbosContext{
+		launched := dbosCtx.launched.Load()
+		// Create a new context that is not canceled when the parent is canceled
+		// but retains all other values
+		childCtx := &dbosContext{
 			ctx:                     context.WithoutCancel(dbosCtx.ctx),
 			logger:                  dbosCtx.logger,
 			systemDB:                dbosCtx.systemDB,
@@ -228,6 +234,8 @@ func WithoutCancel(ctx DBOSContext) DBOSContext {
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
 		}
+		childCtx.launched.Store(launched)
+		return childCtx
 	}
 	return nil
 }
@@ -240,8 +248,9 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 		return nil, func() {}
 	}
 	if dbosCtx, ok := ctx.(*dbosContext); ok {
+		launched := dbosCtx.launched.Load()
 		newCtx, cancelFunc := context.WithTimeoutCause(dbosCtx.ctx, timeout, errors.New("DBOS context timeout"))
-		return &dbosContext{
+		childCtx := &dbosContext{
 			ctx:                     newCtx,
 			logger:                  dbosCtx.logger,
 			systemDB:                dbosCtx.systemDB,
@@ -251,7 +260,9 @@ func WithTimeout(ctx DBOSContext, timeout time.Duration) (DBOSContext, context.C
 			applicationVersion:      dbosCtx.applicationVersion,
 			executorID:              dbosCtx.executorID,
 			applicationID:           dbosCtx.applicationID,
-		}, cancelFunc
+		}
+		childCtx.launched.Store(launched)
+		return childCtx, cancelFunc
 	}
 	return nil, func() {}
 }
@@ -275,6 +286,34 @@ func (c *dbosContext) GetApplicationID() string {
 	return c.applicationID
 }
 
+// ListRegisteredWorkflows returns information about registered workflows with their registration parameters.
+// Supports filtering using functional options.
+func (c *dbosContext) ListRegisteredWorkflows(_ DBOSContext, opts ...ListRegisteredWorkflowsOption) ([]WorkflowRegistryEntry, error) {
+	// Initialize parameters with defaults
+	params := &listRegisteredWorkflowsOptions{}
+
+	// Apply all provided options
+	for _, opt := range opts {
+		opt(params)
+	}
+
+	// Get all registered workflows and apply filters
+	var filteredWorkflows []WorkflowRegistryEntry
+	c.workflowRegistry.Range(func(key, value interface{}) bool {
+		workflow := value.(WorkflowRegistryEntry)
+
+		// Filter by scheduled only
+		if params.scheduledOnly && workflow.CronSchedule == "" {
+			return true
+		}
+
+		filteredWorkflows = append(filteredWorkflows, workflow)
+		return true
+	})
+
+	return filteredWorkflows, nil
+}
+
 // NewDBOSContext creates a new DBOS context with the provided configuration.
 // The context must be launched with Launch() for workflow execution and should be shut down with Shutdown().
 // This function initializes the DBOS system database, sets up the queue sub-system, and prepares the workflow registry.
@@ -317,15 +356,14 @@ func NewDBOSContext(ctx context.Context, inputConfig Config) (DBOSContext, error
 
 	// Register types we serialize with gob
 	var t time.Time
-	gob.Register(t)
+	safeGobRegister(t, initExecutor.logger)
 	var ws []WorkflowStatus
-	gob.Register(ws)
+	safeGobRegister(ws, initExecutor.logger)
 	var si []StepInfo
-	gob.Register(si)
 	var h workflowHandle[any]
-	gob.Register(h)
+	safeGobRegister(h, initExecutor.logger)
 	var ph workflowPollingHandle[any]
-	gob.Register(ph)
+	safeGobRegister(ph, initExecutor.logger)
 
 	// Initialize global variables from processed config (already handles env vars and defaults)
 	initExecutor.applicationVersion = config.ApplicationVersion

dbos/queue.go

@@ -221,7 +221,7 @@ func (qr *queueRunner) run(ctx *dbosContext) {
 					qr.logger.Error("workflow function not found in registry", "workflow_name", workflow.name)
 					continue
 				}
-				registeredWorkflow, ok := registeredWorkflowAny.(workflowRegistryEntry)
+				registeredWorkflow, ok := registeredWorkflowAny.(WorkflowRegistryEntry)
 				if !ok {
 					qr.logger.Error("invalid workflow registry entry type", "workflow_name", workflow.name)
 					continue

dbos/recovery.go

@@ -48,7 +48,7 @@ func recoverPendingWorkflows(ctx *dbosContext, executorIDs []string) ([]Workflow
 			ctx.logger.Error("Workflow function not found in registry", "workflow_id", workflow.ID, "name", workflow.Name)
 			continue
 		}
-		registeredWorkflow, ok := registeredWorkflowAny.(workflowRegistryEntry)
+		registeredWorkflow, ok := registeredWorkflowAny.(WorkflowRegistryEntry)
 		if !ok {
 			ctx.logger.Error("invalid workflow registry entry type", "workflow_id", workflow.ID, "name", workflow.Name)
 			continue