consolidate stub types

beam-cloud · Dec 31, 2024 · 8d5cb47 · 8d5cb47
1 parent e39639c
commit 8d5cb47
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 45 deletions.
diff --git a/pkg/abstractions/shell/shell.go b/pkg/abstractions/shell/shell.go
@@ -12,7 +12,6 @@ import (
 	"github.com/labstack/echo/v4"
 
 	abstractions "github.com/beam-cloud/beta9/pkg/abstractions/common"
-	"github.com/beam-cloud/beta9/pkg/task"
 	pb "github.com/beam-cloud/beta9/proto"
 
 	"github.com/beam-cloud/beta9/pkg/auth"
@@ -46,23 +45,19 @@ type SSHShellService struct {
 	workspaceRepo   repository.WorkspaceRepository
 	containerRepo   repository.ContainerRepository
 	eventRepo       repository.EventRepository
-	taskRepo        repository.TaskRepository
 	tailscale       *network.Tailscale
-	taskDispatcher  *task.Dispatcher
 }
 
 type ShellServiceOpts struct {
-	Config         types.AppConfig
-	RedisClient    *common.RedisClient
-	BackendRepo    repository.BackendRepository
-	WorkspaceRepo  repository.WorkspaceRepository
-	TaskRepo       repository.TaskRepository
-	ContainerRepo  repository.ContainerRepository
-	Scheduler      *scheduler.Scheduler
-	RouteGroup     *echo.Group
-	Tailscale      *network.Tailscale
-	TaskDispatcher *task.Dispatcher
-	EventRepo      repository.EventRepository
+	Config        types.AppConfig
+	RedisClient   *common.RedisClient
+	BackendRepo   repository.BackendRepository
+	WorkspaceRepo repository.WorkspaceRepository
+	ContainerRepo repository.ContainerRepository
+	Scheduler     *scheduler.Scheduler
+	RouteGroup    *echo.Group
+	Tailscale     *network.Tailscale
+	EventRepo     repository.EventRepository
 }
 
 func NewSSHShellService(
@@ -83,9 +78,7 @@ func NewSSHShellService(
 		backendRepo:     opts.BackendRepo,
 		workspaceRepo:   opts.WorkspaceRepo,
 		containerRepo:   opts.ContainerRepo,
-		taskRepo:        opts.TaskRepo,
 		tailscale:       opts.Tailscale,
-		taskDispatcher:  opts.TaskDispatcher,
 		eventRepo:       opts.EventRepo,
 	}
 

diff --git a/pkg/gateway/gateway.go b/pkg/gateway/gateway.go
@@ -379,17 +379,15 @@ func (g *Gateway) registerServices() error {
 
 	// Register shell service
 	ss, err := _shell.NewSSHShellService(g.ctx, _shell.ShellServiceOpts{
-		Config:         g.Config,
-		RedisClient:    g.RedisClient,
-		Scheduler:      g.Scheduler,
-		BackendRepo:    g.BackendRepo,
-		WorkspaceRepo:  g.WorkspaceRepo,
-		ContainerRepo:  g.ContainerRepo,
-		TaskRepo:       g.TaskRepo,
-		Tailscale:      g.Tailscale,
-		TaskDispatcher: g.TaskDispatcher,
-		EventRepo:      g.EventRepo,
-		RouteGroup:     g.rootRouteGroup,
+		Config:        g.Config,
+		RedisClient:   g.RedisClient,
+		Scheduler:     g.Scheduler,
+		BackendRepo:   g.BackendRepo,
+		WorkspaceRepo: g.WorkspaceRepo,
+		ContainerRepo: g.ContainerRepo,
+		Tailscale:     g.Tailscale,
+		EventRepo:     g.EventRepo,
+		RouteGroup:    g.rootRouteGroup,
 	})
 	if err != nil {
 		return err

diff --git a/pkg/gateway/services/stub.go b/pkg/gateway/services/stub.go
@@ -262,7 +262,7 @@ func (gws *GatewayService) GetURL(ctx context.Context, in *pb.GetURLRequest) (*p
 			Ok:  true,
 			Url: invokeUrl,
 		}, nil
-	} else if stub.Type.IsShell() {
+	} else if stub.Type.Kind() == types.StubTypeShell {
 		invokeUrl := common.BuildShellURL(gws.appConfig.GatewayService.HTTP.GetExternalURL(), in.UrlType, stub)
 		return &pb.GetURLResponse{
 			Ok:  true,

diff --git a/pkg/repository/backend_postgres_migrations/020_add_shell_stub_type.go b/pkg/repository/backend_postgres_migrations/020_add_shell_stub_type.go
@@ -13,7 +13,7 @@ func init() {
 }
 
 func upAddShellStubType(ctx context.Context, tx *sql.Tx) error {
-	newStubTypes := []string{"schedule/shell", "asgi/shell", "endpoint/shell", "taskqueue/shell", "function/shell"}
+	newStubTypes := []string{"shell"}
 
 	for _, stubType := range newStubTypes {
 		addEnumSQL := fmt.Sprintf(`

diff --git a/pkg/types/backend.go b/pkg/types/backend.go
@@ -207,20 +207,17 @@ const (
 	StubTypeFunction               string = "function"
 	StubTypeFunctionDeployment     string = "function/deployment"
 	StubTypeFunctionServe          string = "function/serve"
-	StubTypeFunctionShell          string = "function/shell"
 	StubTypeContainer              string = "container"
+	StubTypeShell                  string = "shell"
 	StubTypeTaskQueue              string = "taskqueue"
 	StubTypeTaskQueueDeployment    string = "taskqueue/deployment"
 	StubTypeTaskQueueServe         string = "taskqueue/serve"
-	StubTypeTaskQueueShell         string = "taskqueue/shell"
 	StubTypeEndpoint               string = "endpoint"
 	StubTypeEndpointDeployment     string = "endpoint/deployment"
 	StubTypeEndpointServe          string = "endpoint/serve"
-	StubTypeEndpointShell          string = "endpoint/shell"
 	StubTypeASGI                   string = "asgi"
 	StubTypeASGIDeployment         string = "asgi/deployment"
 	StubTypeASGIServe              string = "asgi/serve"
-	StubTypeASGIShell              string = "asgi/shell"
 	StubTypeScheduledJob           string = "schedule"
 	StubTypeScheduledJobDeployment string = "schedule/deployment"
 	StubTypeScheduledJobShell      string = "schedule/shell"
@@ -231,10 +228,6 @@ const (
 
 type StubType string
 
-func (t StubType) IsShell() bool {
-	return strings.HasSuffix(string(t), "/shell")
-}
-
 func (t StubType) IsServe() bool {
 	return strings.HasSuffix(string(t), "/serve")
 }

diff --git a/sdk/src/beta9/abstractions/base/runner.py b/sdk/src/beta9/abstractions/base/runner.py
@@ -46,12 +46,7 @@
 ASGI_STUB_TYPE = "asgi"
 SCHEDULE_STUB_TYPE = "schedule"
 BOT_STUB_TYPE = "bot"
-
-TASKQUEUE_SHELL_STUB_TYPE = "taskqueue/shell"
-FUNCTION_SHELL_STUB_TYPE = "function/shell"
-ENDPOINT_SHELL_STUB_TYPE = "endpoint/shell"
-ASGI_SHELL_STUB_TYPE = "asgi/shell"
-SCHEDULE_SHELL_STUB_TYPE = "schedule/shell"
+SHELL_STUB_TYPE = "shell"
 
 TASKQUEUE_DEPLOYMENT_STUB_TYPE = "taskqueue/deployment"
 ENDPOINT_DEPLOYMENT_STUB_TYPE = "endpoint/deployment"

diff --git a/sdk/src/beta9/abstractions/endpoint.py b/sdk/src/beta9/abstractions/endpoint.py
@@ -15,8 +15,8 @@
     ASGI_STUB_TYPE,
     ENDPOINT_DEPLOYMENT_STUB_TYPE,
     ENDPOINT_SERVE_STUB_TYPE,
-    ENDPOINT_SHELL_STUB_TYPE,
     ENDPOINT_STUB_TYPE,
+    SHELL_STUB_TYPE,
     RunnerAbstraction,
 )
 from ..abstractions.image import Image
@@ -614,7 +614,7 @@ def notify(*_, **__):
 
     @with_grpc_error_handling
     def shell(self, timeout: int = 0, url_type: str = ""):
-        stub_type = ENDPOINT_SHELL_STUB_TYPE
+        stub_type = SHELL_STUB_TYPE
 
         if not self.parent.prepare_runtime(
             func=self.func, stub_type=stub_type, force_create_stub=True