fix: resolve the pg DSN from runtime in the proxy

TBD54566975 · Nov 22, 2024 · f2a91ef · f2a91ef
1 parent f5554a1
commit f2a91ef
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/backend/controller/controller.go b/backend/controller/controller.go
@@ -769,9 +769,21 @@ func (s *Service) GetModuleContext(ctx context.Context, req *connect.Request[ftl
 						continue
 					}
 					dbTypes[db.Name] = dbType
+					// TODO: Move the DSN resolution to the runtime
 					if db.Runtime != nil {
+						var dsn string
+						switch dbType {
+						case modulecontext.DBTypePostgres:
+							// TODO: Get the port from config
+							dsn = "postgres://127.0.0.1:5678/" + db.Name
+						case modulecontext.DBTypeMySQL:
+							// TODO: Route MySQL through a proxy as well
+							dsn = db.Runtime.DSN
+						default:
+							return connect.NewError(connect.CodeInternal, fmt.Errorf("unknown DB type: %s", db.Type))
+						}
 						databases[db.Name] = modulecontext.Database{
-							DSN:    db.Runtime.DSN,
+							DSN:    dsn,
 							DBType: dbType,
 						}
 					}

diff --git a/backend/runner/runner.go b/backend/runner/runner.go
@@ -162,7 +162,7 @@ func Start(ctx context.Context, config Config) error {
 
 	g, ctx := errgroup.WithContext(ctx)
 	g.Go(func() error {
-		return svc.StartPgProxy(ctx, module)
+		return svc.startPgProxy(ctx, module)
 	})
 	g.Go(func() error {
 		// TODO: Make sure pgproxy is ready before starting the runner
@@ -593,11 +593,20 @@ func (s *Service) healthCheck(writer http.ResponseWriter, request *http.Request)
 	writer.WriteHeader(http.StatusServiceUnavailable)
 }
 
-func (s *Service) StartPgProxy(ctx context.Context, module *schema.Module) error {
-	if err := pgproxy.New(s.config.PgProxyConfig, func(ctx context.Context, params map[string]string) (string, error) {
-		db := params["database"]
+func (s *Service) startPgProxy(ctx context.Context, module *schema.Module) error {
+	databases := map[string]*schema.Database{}
+	for _, decl := range module.Decls {
+		if db, ok := decl.(*schema.Database); ok {
+			databases[db.Name] = db
+		}
+	}
 
-		return "postgres://127.0.0.1:5432/" + db + "?user=" + params["user"], nil
+	if err := pgproxy.New(s.config.PgProxyConfig, func(ctx context.Context, params map[string]string) (string, error) {
+		db, ok := databases[params["database"]]
+		if !ok {
+			return "", fmt.Errorf("database %s not found", params["database"])
+		}
+		return db.Runtime.DSN, nil
 	}).Start(ctx); err != nil {
 		return fmt.Errorf("failed to start pgproxy: %w", err)
 	}