allow custom workflow names

Author: maxdml

dbos/workflow.go

@@ -205,10 +205,11 @@ type WrappedWorkflowFunc func(ctx DBOSContext, input any, opts ...WorkflowOption
 type workflowRegistryEntry struct {
 	wrappedFunction WrappedWorkflowFunc
 	maxRetries      int
+	name            string
 }
 
 // Register adds a workflow function to the registry (thread-safe, only once per name)
-func registerWorkflow(ctx DBOSContext, workflowName string, fn WrappedWorkflowFunc, maxRetries int) {
+func registerWorkflow(ctx DBOSContext, workflowFQN string, fn WrappedWorkflowFunc, maxRetries int, customName string) {
 	// Skip if we don't have a concrete dbosContext
 	c, ok := ctx.(*dbosContext)
 	if !ok {
@@ -222,14 +223,15 @@ func registerWorkflow(ctx DBOSContext, workflowName string, fn WrappedWorkflowFu
 	c.workflowRegMutex.Lock()
 	defer c.workflowRegMutex.Unlock()
 
-	if _, exists := c.workflowRegistry[workflowName]; exists {
-		c.logger.Error("workflow function already registered", "fqn", workflowName)
-		panic(newConflictingRegistrationError(workflowName))
+	if _, exists := c.workflowRegistry[workflowFQN]; exists {
+		c.logger.Error("workflow function already registered", "fqn", workflowFQN)
+		panic(newConflictingRegistrationError(workflowFQN))
 	}
 
-	c.workflowRegistry[workflowName] = workflowRegistryEntry{
+	c.workflowRegistry[workflowFQN] = workflowRegistryEntry{
 		wrappedFunction: fn,
 		maxRetries:      maxRetries,
+		name:            customName,
 	}
 }
 
@@ -275,6 +277,7 @@ func registerScheduledWorkflow(ctx DBOSContext, workflowName string, fn Workflow
 type workflowRegistrationParams struct {
 	cronSchedule string
 	maxRetries   int
+	name         string
 }
 
 type workflowRegistrationOption func(*workflowRegistrationParams)
@@ -295,6 +298,12 @@ func WithSchedule(schedule string) workflowRegistrationOption {
 	}
 }
 
+func WithWorkflowName(name string) workflowRegistrationOption {
+	return func(p *workflowRegistrationParams) {
+		p.name = name
+	}
+}
+
 // RegisterWorkflow registers the provided function as a durable workflow with the provided DBOSContext workflow registry
 // If the workflow is a scheduled workflow (determined by the presence of a cron schedule), it will also register a cron job to execute it
 // RegisterWorkflow is generically typed, providing compile-time type checking and allowing us to register the workflow input and output types for gob encoding
@@ -348,7 +357,7 @@ func RegisterWorkflow[P any, R any](ctx DBOSContext, fn GenericWorkflowFunc[P, R
 		}
 		return &workflowPollingHandle[any]{workflowID: handle.GetWorkflowID(), dbosContext: ctx}, nil // this is only used by recovery and queue runner so far -- queue runner dismisses it
 	})
-	registerWorkflow(ctx, fqn, typeErasedWrapper, registrationParams.maxRetries)
+	registerWorkflow(ctx, fqn, typeErasedWrapper, registrationParams.maxRetries, registrationParams.name)
 
 	// If this is a scheduled workflow, register a cron job
 	if registrationParams.cronSchedule != "" {
@@ -398,6 +407,7 @@ func WithApplicationVersion(version string) WorkflowOption {
 	}
 }
 
+// An internal option we use to map the reflection function name to the registration options.
 func withWorkflowName(name string) WorkflowOption {
 	return func(p *workflowParams) {
 		p.workflowName = name
@@ -486,6 +496,11 @@ func (c *dbosContext) RunAsWorkflow(_ DBOSContext, fn WorkflowFunc, input any, o
 		params.maxRetries = registeredWorkflow.maxRetries
 	}
 
+	// Use the custom workflow name if it was provided during registration
+	if registeredWorkflow.name != "" {
+		params.workflowName = registeredWorkflow.name
+	}
+
 	// Check if we are within a workflow (and thus a child workflow)
 	parentWorkflowState, ok := c.Value(workflowStateKey).(*workflowState)
 	isChildWorkflow := ok && parentWorkflowState != nil

dbos/workflows_test.go

@@ -300,6 +300,74 @@ func TestWorkflowsRegistration(t *testing.T) {
 			}
 		})
 	}
+
+	t.Run("DoubleRegistrationWithoutName", func(t *testing.T) {
+		// Create a fresh DBOS context for this test
+		freshCtx := setupDBOS(t, false, false) // Don't check for leaks and don't reset DB
+
+		// First registration should work
+		RegisterWorkflow(freshCtx, simpleWorkflow)
+
+		// Second registration of the same workflow should panic with ConflictingRegistrationError
+		defer func() {
+			r := recover()
+			if r == nil {
+				t.Fatal("expected panic from double registration but got none")
+			}
+			dbosErr, ok := r.(*DBOSError)
+			if !ok {
+				t.Fatalf("expected panic to be *DBOSError, got %T", r)
+			}
+			if dbosErr.Code != ConflictingRegistrationError {
+				t.Fatalf("expected ConflictingRegistrationError, got %v", dbosErr.Code)
+			}
+		}()
+		RegisterWorkflow(freshCtx, simpleWorkflow)
+	})
+
+	t.Run("DoubleRegistrationWithCustomName", func(t *testing.T) {
+		// Create a fresh DBOS context for this test
+		freshCtx := setupDBOS(t, false, false) // Don't check for leaks and don't reset DB
+
+		// First registration with custom name should work
+		RegisterWorkflow(freshCtx, simpleWorkflow, WithWorkflowName("custom-workflow"))
+
+		// Second registration with same custom name should panic with ConflictingRegistrationError
+		defer func() {
+			r := recover()
+			if r == nil {
+				t.Fatal("expected panic from double registration with custom name but got none")
+			}
+			dbosErr, ok := r.(*DBOSError)
+			if !ok {
+				t.Fatalf("expected panic to be *DBOSError, got %T", r)
+			}
+			if dbosErr.Code != ConflictingRegistrationError {
+				t.Fatalf("expected ConflictingRegistrationError, got %v", dbosErr.Code)
+			}
+		}()
+		RegisterWorkflow(freshCtx, simpleWorkflow, WithWorkflowName("custom-workflow"))
+	})
+
+	t.Run("RegisterAfterLaunchPanics", func(t *testing.T) {
+		// Create a fresh DBOS context for this test
+		freshCtx := setupDBOS(t, false, false) // Don't check for leaks and don't reset DB
+
+		// Launch DBOS context
+		err := freshCtx.Launch()
+		if err != nil {
+			t.Fatalf("failed to launch DBOS context: %v", err)
+		}
+		defer freshCtx.Cancel()
+
+		// Attempting to register after launch should panic
+		defer func() {
+			if r := recover(); r == nil {
+				t.Fatal("expected panic from registration after launch but got none")
+			}
+		}()
+		RegisterWorkflow(freshCtx, simpleWorkflow)
+	})
 }
 
 func stepWithinAStep(ctx context.Context) (string, error) {