fix: module context should be immutable

block · May 3, 2024 · cdbecf8 · cdbecf8
1 parent 519f2fe
commit cdbecf8
Show file tree

Hide file tree

Showing 11 changed files with 181 additions and 162 deletions.
diff --git a/backend/controller/controller.go b/backend/controller/controller.go
@@ -640,11 +640,11 @@ func (s *Service) GetModuleContext(ctx context.Context, req *connect.Request[ftl
 	if !ok {
 		return nil, connect.NewError(connect.CodeNotFound, fmt.Errorf("module %q not found", req.Msg.Module))
 	}
-	moduleContext, err := modulecontext.FromEnvironment(ctx, module.Name, false)
+	moduleContext, err := modulecontext.New(module.Name).UpdateFromEnvironment(ctx)
 	if err != nil {
 		return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("could not get module context: %w", err))
 	}
-	response, err := moduleContext.ToProto(module.Name)
+	response, err := moduleContext.ToProto()
 	if err != nil {
 		return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("could not marshal module context: %w", err))
 	}

diff --git a/go-runtime/ftl/call.go b/go-runtime/ftl/call.go
@@ -9,13 +9,14 @@ import (
 
 	ftlv1 "github.com/TBD54566975/ftl/backend/protos/xyz/block/ftl/v1"
 	"github.com/TBD54566975/ftl/backend/protos/xyz/block/ftl/v1/ftlv1connect"
+	"github.com/TBD54566975/ftl/backend/schema"
 	"github.com/TBD54566975/ftl/go-runtime/encoding"
 	"github.com/TBD54566975/ftl/go-runtime/modulecontext"
 	"github.com/TBD54566975/ftl/internal/rpc"
 )
 
 func call[Req, Resp any](ctx context.Context, callee Ref, req Req, inline Verb[Req, Resp]) (resp Resp, err error) {
-	behavior, err := modulecontext.FromContext(ctx).BehaviorForVerb(modulecontext.Ref(callee))
+	behavior, err := modulecontext.FromContext(ctx).BehaviorForVerb(schema.Ref{Module: callee.Module, Name: callee.Name})
 	if err != nil {
 		return resp, fmt.Errorf("%s: %w", callee, err)
 	}

diff --git a/go-runtime/ftl/config_test.go b/go-runtime/ftl/config_test.go
@@ -2,6 +2,7 @@ package ftl
 
 import (
 	"context"
+	"encoding/json"
 	"testing"
 
 	"github.com/alecthomas/assert/v2"
@@ -11,16 +12,25 @@ import (
 )
 
 func TestConfig(t *testing.T) {
-	ctx := log.ContextWithNewDefaultLogger(context.Background())
-
-	moduleCtx := modulecontext.New("test")
-	ctx = moduleCtx.ApplyToContext(ctx)
-
 	type C struct {
 		One string
 		Two string
 	}
+
+	ctx := log.ContextWithNewDefaultLogger(context.Background())
+
+	data, err := json.Marshal(C{"one", "two"})
+	assert.NoError(t, err)
+
+	moduleCtx := modulecontext.New("test").Update(
+		map[string][]byte{
+			"test": data,
+		},
+		map[string][]byte{},
+		map[string]modulecontext.Database{},
+	)
+	ctx = moduleCtx.ApplyToContext(ctx)
+
 	config := Config[C]("test")
-	assert.NoError(t, moduleCtx.SetConfig("test", C{"one", "two"}))
 	assert.Equal(t, C{"one", "two"}, config.Get(ctx))
 }
diff --git a/go-runtime/ftl/ftltest/ftltest.go b/go-runtime/ftl/ftltest/ftltest.go
@@ -3,30 +3,47 @@ package ftltest
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"reflect"
 
+	"github.com/TBD54566975/ftl/backend/schema"
 	"github.com/TBD54566975/ftl/go-runtime/ftl"
 	"github.com/TBD54566975/ftl/go-runtime/modulecontext"
 	"github.com/TBD54566975/ftl/internal/log"
 )
 
+type Options struct {
+	configs                 map[string][]byte
+	secrets                 map[string][]byte
+	mockVerbs               map[schema.RefKey]modulecontext.MockVerb
+	allowDirectVerbBehavior bool
+}
+
 // Context suitable for use in testing FTL verbs with provided options
-func Context(options ...func(context.Context) error) context.Context {
+func Context(options ...func(*Options) error) context.Context {
 	ctx := log.ContextWithNewDefaultLogger(context.Background())
-	context, err := modulecontext.FromEnvironment(ctx, ftl.Module(), true)
-	if err != nil {
-		panic(err)
-	}
-	ctx = context.ApplyToContext(ctx)
 
+	state := &Options{
+		configs:   make(map[string][]byte),
+		secrets:   make(map[string][]byte),
+		mockVerbs: make(map[schema.RefKey]modulecontext.MockVerb),
+	}
 	for _, option := range options {
-		err = option(ctx)
+		err := option(state)
 		if err != nil {
 			panic(fmt.Sprintf("error applying option: %v", err))
 		}
 	}
-	return ctx
+
+	moduleCtx := modulecontext.New(ftl.Module())
+	moduleCtx, err := moduleCtx.UpdateFromEnvironment(ctx)
+	if err != nil {
+		panic(fmt.Sprintf("error setting up module context from environment: %v", err))
+	}
+	moduleCtx = moduleCtx.Update(state.configs, state.secrets, map[string]modulecontext.Database{})
+	moduleCtx = moduleCtx.UpdateForTesting(state.mockVerbs, state.allowDirectVerbBehavior)
+	return moduleCtx.ApplyToContext(ctx)
 }
 
 // WithConfig sets a configuration for the current module
@@ -38,12 +55,17 @@ func Context(options ...func(context.Context) error) context.Context {
 //	... other options
 //
 // )
-func WithConfig[T ftl.ConfigType](config ftl.ConfigValue[T], value T) func(context.Context) error {
-	return func(ctx context.Context) error {
+func WithConfig[T ftl.ConfigType](config ftl.ConfigValue[T], value T) func(*Options) error {
+	return func(state *Options) error {
 		if config.Module != ftl.Module() {
 			return fmt.Errorf("config %v does not match current module %s", config.Module, ftl.Module())
 		}
-		return modulecontext.FromContext(ctx).SetConfig(config.Name, value)
+		data, err := json.Marshal(value)
+		if err != nil {
+			return err
+		}
+		state.configs[config.Name] = data
+		return nil
 	}
 }
 
@@ -56,12 +78,17 @@ func WithConfig[T ftl.ConfigType](config ftl.ConfigValue[T], value T) func(conte
 //	... other options
 //
 // )
-func WithSecret[T ftl.SecretType](secret ftl.SecretValue[T], value T) func(context.Context) error {
-	return func(ctx context.Context) error {
+func WithSecret[T ftl.SecretType](secret ftl.SecretValue[T], value T) func(*Options) error {
+	return func(state *Options) error {
 		if secret.Module != ftl.Module() {
 			return fmt.Errorf("secret %v does not match current module %s", secret.Module, ftl.Module())
 		}
-		return modulecontext.FromContext(ctx).SetSecret(secret.Name, value)
+		data, err := json.Marshal(value)
+		if err != nil {
+			return err
+		}
+		state.secrets[secret.Name] = data
+		return nil
 	}
 }
 
@@ -70,32 +97,32 @@ func WithSecret[T ftl.SecretType](secret ftl.SecretValue[T], value T) func(conte
 // To be used when setting up a context for a test:
 // ctx := ftltest.Context(
 //
-//	ftltest.WhenVerb(Example.Verb, func(ctx context.Context, req Example.Req) (Example.Resp, error) {
+//	ftltest.WhenVerb(Example.Verb, func(state *OptionsState, req Example.Req) (Example.Resp, error) {
 //	    ...
 //	}),
 //	... other options
 //
 // )
-func WhenVerb[Req any, Resp any](verb ftl.Verb[Req, Resp], fake func(ctx context.Context, req Req) (resp Resp, err error)) func(context.Context) error {
-	return func(ctx context.Context) error {
+func WhenVerb[Req any, Resp any](verb ftl.Verb[Req, Resp], fake func(ctx context.Context, req Req) (resp Resp, err error)) func(*Options) error {
+	return func(state *Options) error {
 		ref := ftl.FuncRef(verb)
-		modulecontext.FromContext(ctx).SetMockVerb(modulecontext.Ref(ref), func(ctx context.Context, req any) (resp any, err error) {
+		state.mockVerbs[schema.RefKey(ref)] = func(ctx context.Context, req any) (resp any, err error) {
 			request, ok := req.(Req)
 			if !ok {
 				return nil, fmt.Errorf("invalid request type %T for %v, expected %v", req, ref, reflect.TypeFor[Req]())
 			}
 			return fake(ctx, request)
-		})
+		}
 		return nil
 	}
 }
 
 // WithCallsAllowedWithinModule allows tests to enable calls to all verbs within the current module
 //
 // Any overrides provided by calling WhenVerb(...) will take precedence
-func WithCallsAllowedWithinModule() func(context.Context) error {
-	return func(ctx context.Context) error {
-		modulecontext.FromContext(ctx).AllowDirectVerbBehaviorWithinModule()
+func WithCallsAllowedWithinModule() func(*Options) error {
+	return func(state *Options) error {
+		state.allowDirectVerbBehavior = true
 		return nil
 	}
 }
diff --git a/go-runtime/ftl/secrets_test.go b/go-runtime/ftl/secrets_test.go
@@ -2,6 +2,7 @@ package ftl
 
 import (
 	"context"
+	"encoding/json"
 	"testing"
 
 	"github.com/alecthomas/assert/v2"
@@ -11,16 +12,25 @@ import (
 )
 
 func TestSecret(t *testing.T) {
-	ctx := log.ContextWithNewDefaultLogger(context.Background())
-
-	moduleCtx := modulecontext.New("test")
-	ctx = moduleCtx.ApplyToContext(ctx)
-
 	type C struct {
 		One string
 		Two string
 	}
+
+	ctx := log.ContextWithNewDefaultLogger(context.Background())
+
+	data, err := json.Marshal(C{"one", "two"})
+	assert.NoError(t, err)
+
+	moduleCtx := modulecontext.New("test").Update(
+		map[string][]byte{},
+		map[string][]byte{
+			"test": data,
+		},
+		map[string]modulecontext.Database{},
+	)
+	ctx = moduleCtx.ApplyToContext(ctx)
+
 	secret := Secret[C]("test")
-	assert.NoError(t, moduleCtx.SetSecret("test", C{"one", "two"}))
 	assert.Equal(t, C{"one", "two"}, secret.Get(ctx))
 }
diff --git a/go-runtime/modulecontext/from_environment.go b/go-runtime/modulecontext/from_environment.go
@@ -9,42 +9,35 @@ import (
 	cf "github.com/TBD54566975/ftl/common/configuration"
 )
 
-// FromEnvironment creates a ModuleContext from the local environment.
+// UpdateFromEnvironment copies a ModuleContext and gathers configs, secrets and databases from the local environment.
 //
 // This is useful for testing and development, where the environment is used to provide
 // configurations, secrets and DSNs. The context is built from a combination of
 // the ftl-project.toml file and (for now) environment variables.
-func FromEnvironment(ctx context.Context, module string, isTesting bool) (*ModuleContext, error) {
+func (m ModuleContext) UpdateFromEnvironment(ctx context.Context) (ModuleContext, error) {
 	// TODO: split this func into separate purposes: explicitly reading a particular project file, and reading DSNs from environment
-	var moduleCtx *ModuleContext
-	if isTesting {
-		moduleCtx = NewForTesting(module)
-	} else {
-		moduleCtx = New(module)
-	}
-
 	cm, err := cf.NewDefaultConfigurationManagerFromEnvironment(ctx)
 	if err != nil {
-		return nil, err
+		return ModuleContext{}, err
 	}
-	configs, err := cm.MapForModule(ctx, module)
+	configs, err := cm.MapForModule(ctx, m.module)
 	if err != nil {
-		return nil, err
+		return ModuleContext{}, err
 	}
 	for name, data := range configs {
-		moduleCtx.SetConfigData(name, data)
+		m.configs[name] = data
 	}
 
 	sm, err := cf.NewDefaultSecretsManagerFromEnvironment(ctx)
 	if err != nil {
-		return nil, err
+		return ModuleContext{}, err
 	}
-	secrets, err := sm.MapForModule(ctx, module)
+	secrets, err := sm.MapForModule(ctx, m.module)
 	if err != nil {
-		return nil, err
+		return ModuleContext{}, err
 	}
 	for name, data := range secrets {
-		moduleCtx.SetSecretData(name, data)
+		m.secrets[name] = data
 	}
 
 	for _, entry := range os.Environ() {
@@ -53,23 +46,25 @@ func FromEnvironment(ctx context.Context, module string, isTesting bool) (*Modul
 		}
 		parts := strings.SplitN(entry, "=", 2)
 		if len(parts) != 2 {
-			return nil, fmt.Errorf("invalid DSN environment variable: %s", entry)
+			return ModuleContext{}, fmt.Errorf("invalid DSN environment variable: %s", entry)
 		}
 		key := parts[0]
 		value := parts[1]
 		// FTL_POSTGRES_DSN_MODULE_DBNAME
 		parts = strings.Split(key, "_")
 		if len(parts) != 5 {
-			return nil, fmt.Errorf("invalid DSN environment variable: %s", entry)
+			return ModuleContext{}, fmt.Errorf("invalid DSN environment variable: %s", entry)
 		}
 		moduleName := parts[3]
 		dbName := parts[4]
-		if !strings.EqualFold(moduleName, module) {
+		if !strings.EqualFold(moduleName, m.module) {
 			continue
 		}
-		if err := moduleCtx.AddDatabase(strings.ToLower(dbName), DBTypePostgres, value); err != nil {
-			return nil, err
+		db, err := NewDatabase(DBTypePostgres, value)
+		if err != nil {
+			return ModuleContext{}, fmt.Errorf("could not create database %q with DSN %q: %w", dbName, value, err)
 		}
+		m.databases[dbName] = db
 	}
-	return moduleCtx, nil
+	return m, nil
 }
diff --git a/go-runtime/modulecontext/from_environment_test.go b/go-runtime/modulecontext/from_environment_test.go
@@ -37,10 +37,10 @@ func TestFromEnvironment(t *testing.T) {
 
 	ctx := log.ContextWithNewDefaultLogger(context.Background())
 
-	moduleContext, err := FromEnvironment(ctx, "echo", false)
+	moduleContext, err := New("echo").UpdateFromEnvironment(ctx)
 	assert.NoError(t, err)
 
-	response, err := moduleContext.ToProto("echo")
+	response, err := moduleContext.ToProto()
 	assert.NoError(t, err)
 
 	assert.Equal(t, &ftlv1.ModuleContextResponse{

diff --git a/go-runtime/modulecontext/from_proto.go b/go-runtime/modulecontext/from_proto.go
@@ -1,21 +1,19 @@
 package modulecontext
 
 import (
+	"fmt"
+
 	ftlv1 "github.com/TBD54566975/ftl/backend/protos/xyz/block/ftl/v1"
 )
 
-func FromProto(response *ftlv1.ModuleContextResponse) (*ModuleContext, error) {
-	moduleCtx := New(response.Module)
-	for name, data := range response.Configs {
-		moduleCtx.configs[name] = data
-	}
-	for name, data := range response.Secrets {
-		moduleCtx.secrets[name] = data
-	}
-	for _, entry := range response.Databases {
-		if err := moduleCtx.AddDatabase(entry.Name, DBType(entry.Type), entry.Dsn); err != nil {
-			return nil, err
+func FromProto(response *ftlv1.ModuleContextResponse) (ModuleContext, error) {
+	databases := map[string]Database{}
+	for name, entry := range response.Databases {
+		db, err := NewDatabase(DBType(entry.Type), entry.Dsn)
+		if err != nil {
+			return ModuleContext{}, fmt.Errorf("could not create database %q with DSN %q: %w", name, entry.Dsn, err)
 		}
+		databases[entry.Name] = db
 	}
-	return moduleCtx, nil
+	return New(response.Module).Update(response.Configs, response.Secrets, databases), nil
 }