diff --git a/pkg/abstractions/shell/shell.go b/pkg/abstractions/shell/shell.go index cb0928c59..d98c0318a 100644 --- 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 index 54817526a..8c5b7f7a2 100644 --- 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 index 22416de4e..fdffdb8ae 100644 --- 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 index c1843c6f2..0747a0705 100644 --- 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 index 19d34d777..aafbfe4aa 100644 --- 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 index a30e0c992..cf5deaeee 100644 --- 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 index 591dfb3de..452c30180 100644 --- 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