Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 21 additions & 14 deletions pkg/cmd/copy/copy.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ type CopyStore interface {

func NewCmdCopy(t *terminal.Terminal, store CopyStore, noLoginStartStore CopyStore) *cobra.Command {
var host bool
var org string
cmd := &cobra.Command{
Annotations: map[string]string{"access": ""},
Use: "copy",
Expand All @@ -52,19 +53,25 @@ func NewCmdCopy(t *terminal.Terminal, store CopyStore, noLoginStartStore CopySto
Args: cmderrors.TransformToValidationError(cobra.ExactArgs(2)),
ValidArgsFunction: completions.GetAllWorkspaceNameCompletionHandler(noLoginStartStore, t),
RunE: func(cmd *cobra.Command, args []string) error {
err := runCopyCommand(t, store, args[0], args[1], host)
err := runCopyCommand(t, store, args[0], args[1], host, org)
if err != nil {
return breverrors.WrapAndTrace(err)
}
return nil
},
}
cmd.Flags().BoolVarP(&host, "host", "", false, "copy to/from the host machine instead of the container")
cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)")
errRegComp := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noLoginStartStore, t))
if errRegComp != nil {
breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(errRegComp))
fmt.Print(breverrors.WrapAndTrace(errRegComp))
}

return cmd
}

func runCopyCommand(t *terminal.Terminal, cstore CopyStore, source, dest string, host bool) error {
func runCopyCommand(t *terminal.Terminal, cstore CopyStore, source, dest string, host bool, orgFlag string) error {
workspaceNameOrID, remotePath, localPath, isUpload, err := parseCopyArguments(source, dest)
if err != nil {
return breverrors.WrapAndTrace(err)
Expand All @@ -77,15 +84,19 @@ func runCopyCommand(t *terminal.Terminal, cstore CopyStore, source, dest string,
}
}

target, err := util.ResolveWorkspaceOrNode(cstore, workspaceNameOrID)
resolvedOrg, err := util.ResolveOrgFromFlag(cstore, orgFlag)
if err != nil {
return breverrors.WrapAndTrace(err)
}
target, err := util.ResolveWorkspaceOrNodeInOrg(cstore, workspaceNameOrID, resolvedOrg.ID)
if err != nil {
return breverrors.WrapAndTrace(err)
}
if target.Node != nil {
return copyExternalNode(t, cstore, target.Node, localPath, remotePath, isUpload)
}

workspace, err := prepareWorkspace(t, cstore, target.Workspace)
workspace, err := prepareWorkspace(t, cstore, target.Workspace, resolvedOrg.ID)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand Down Expand Up @@ -126,11 +137,11 @@ func parseCopyArguments(source, dest string) (workspaceNameOrID, remotePath, loc
return destWorkspace, destPath, source, true, nil
}

func prepareWorkspace(t *terminal.Terminal, cstore CopyStore, workspace *entity.Workspace) (*entity.Workspace, error) {
func prepareWorkspace(t *terminal.Terminal, cstore CopyStore, workspace *entity.Workspace, orgID string) (*entity.Workspace, error) {
s := t.NewSpinner()

if workspace.Status == "STOPPED" {
err := startWorkspaceIfStopped(t, s, cstore, workspace.Name, workspace)
err := startWorkspaceIfStopped(t, s, cstore, workspace.Name, workspace, orgID)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
}
Expand All @@ -141,7 +152,7 @@ func prepareWorkspace(t *terminal.Terminal, cstore CopyStore, workspace *entity.
return nil, breverrors.WrapAndTrace(err)
}

workspace, err = util.GetUserWorkspaceByNameOrIDErr(cstore, workspace.Name)
workspace, err = util.GetUserWorkspaceByNameOrIDErrInOrg(cstore, workspace.Name, orgID)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
}
Expand Down Expand Up @@ -268,12 +279,8 @@ func waitForSSHToBeAvailable(sshAlias string, s *spinner.Spinner) error {
}
}

func startWorkspaceIfStopped(t *terminal.Terminal, s *spinner.Spinner, tstore CopyStore, wsIDOrName string, workspace *entity.Workspace) error {
activeOrg, err := tstore.GetActiveOrganizationOrDefault()
if err != nil {
return breverrors.WrapAndTrace(err)
}
workspaces, err := tstore.GetWorkspaceByNameOrID(activeOrg.ID, wsIDOrName)
func startWorkspaceIfStopped(t *terminal.Terminal, s *spinner.Spinner, tstore CopyStore, wsIDOrName string, workspace *entity.Workspace, orgID string) error {
workspaces, err := tstore.GetWorkspaceByNameOrID(orgID, wsIDOrName)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand All @@ -286,7 +293,7 @@ func startWorkspaceIfStopped(t *terminal.Terminal, s *spinner.Spinner, tstore Co
if err != nil {
return breverrors.WrapAndTrace(err)
}
workspace, err = util.GetUserWorkspaceByNameOrIDErr(tstore, wsIDOrName)
workspace, err = util.GetUserWorkspaceByNameOrIDErrInOrg(tstore, wsIDOrName, orgID)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand Down
18 changes: 15 additions & 3 deletions pkg/cmd/exec/exec.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ type ExecStore interface {

func NewCmdExec(t *terminal.Terminal, store ExecStore, noLoginStartStore ExecStore) *cobra.Command {
var host bool
var org string
cmd := &cobra.Command{
Annotations: map[string]string{"access": ""},
Use: "exec [instance...] <command>",
Expand Down Expand Up @@ -87,13 +88,18 @@ func NewCmdExec(t *terminal.Terminal, store ExecStore, noLoginStartStore ExecSto
return breverrors.NewValidationError("command is required")
}

resolvedOrg, err := util.ResolveOrgFromFlag(store, org)
if err != nil {
return breverrors.WrapAndTrace(err)
}

// Run on each instance
var errors error
for _, instanceName := range instanceNames {
if len(instanceNames) > 1 {
fmt.Fprintf(os.Stderr, "\n=== %s ===\n", instanceName)
}
err = runExecCommand(t, store, instanceName, host, cmdToRun)
err = runExecCommand(t, store, instanceName, host, cmdToRun, resolvedOrg.ID)
if err != nil {
if len(instanceNames) > 1 {
fmt.Fprintf(os.Stderr, "Error on %s: %v\n", instanceName, err)
Expand All @@ -114,6 +120,12 @@ func NewCmdExec(t *terminal.Terminal, store ExecStore, noLoginStartStore ExecSto
},
}
cmd.Flags().BoolVarP(&host, "host", "", false, "ssh into the host machine instead of the container")
cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)")
errRegComp := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noLoginStartStore, t))
if errRegComp != nil {
breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(errRegComp))
fmt.Print(breverrors.WrapAndTrace(errRegComp))
}

return cmd
}
Expand Down Expand Up @@ -175,7 +187,7 @@ func parseCommand(command string) (string, error) {

const pollTimeout = 10 * time.Minute

func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID string, host bool, command string) error {
func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID string, host bool, command string, orgID string) error {
// Determine SSH alias: use the workspace name directly (with -host suffix if needed)
sshName := workspaceNameOrID
if host {
Expand All @@ -196,7 +208,7 @@ func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID st
// SSH failed — now check what's going on with the instance
fmt.Fprintf(os.Stderr, "Connection failed, checking instance status...\n")

workspace, lookupErr := util.GetUserWorkspaceByNameOrIDErr(sstore, workspaceNameOrID)
workspace, lookupErr := util.GetUserWorkspaceByNameOrIDErrInOrg(sstore, workspaceNameOrID, orgID)
if lookupErr != nil {
return breverrors.WrapAndTrace(fmt.Errorf(
"ssh connection failed and could not look up instance %q: %w\nPlease check your instances with: brev ls",
Expand Down
16 changes: 8 additions & 8 deletions pkg/cmd/gpucreate/gpucreate.go
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ type GPUCreateStore interface {
util.GetWorkspaceByNameOrIDErrStore
gpusearch.GPUSearchStore
GetActiveOrganizationOrDefault() (*entity.Organization, error)
GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error)
GetCurrentUser() (*entity.User, error)
GetWorkspace(workspaceID string) (*entity.Workspace, error)
CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error)
Expand Down Expand Up @@ -154,6 +155,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
var containerImage string
var composeFile string
var filters searchFilterFlags
var org string

cmd := &cobra.Command{
Annotations: map[string]string{"workspace": ""},
Expand Down Expand Up @@ -230,7 +232,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
ComposeFile: composeFile,
}

err = RunGPUCreate(t, gpuCreateStore, opts)
err = RunGPUCreate(t, gpuCreateStore, opts, org)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand All @@ -239,6 +241,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra
}

registerCreateFlags(cmd, &name, &instanceTypes, &count, &parallel, &detached, &timeout, &startupScript, &dryRun, &mode, &jupyter, &containerImage, &composeFile, &filters)
cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)")

return cmd
}
Expand Down Expand Up @@ -551,7 +554,7 @@ type createContext struct {
}

// newCreateContext initializes the context for instance creation
func newCreateContext(t *terminal.Terminal, store GPUCreateStore, opts GPUCreateOptions) (*createContext, error) {
func newCreateContext(t *terminal.Terminal, store GPUCreateStore, opts GPUCreateOptions, orgFlag string) (*createContext, error) {
piped := gpusearch.IsStdoutPiped()

ctx := &createContext{
Expand All @@ -578,13 +581,10 @@ func newCreateContext(t *terminal.Terminal, store GPUCreateStore, opts GPUCreate
ctx.user = user

// Get organization
org, err := store.GetActiveOrganizationOrDefault()
org, err := util.ResolveOrgFromFlag(store, orgFlag)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
}
if org == nil {
return nil, breverrors.NewValidationError("no organization found")
}
ctx.org = org

// Fetch instance types with workspace groups
Expand Down Expand Up @@ -754,8 +754,8 @@ func (c *createContext) printSummary(workspaces []*entity.Workspace) {
}

// RunGPUCreate executes the GPU create with retry logic
func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUCreateOptions) error {
ctx, err := newCreateContext(t, gpuCreateStore, opts)
func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUCreateOptions, orgFlag string) error {
ctx, err := newCreateContext(t, gpuCreateStore, opts, orgFlag)
if err != nil {
return err
}
Expand Down
7 changes: 7 additions & 0 deletions pkg/cmd/gpucreate/gpucreate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,13 @@ func (m *MockGPUCreateStore) GetActiveOrganizationOrDefault() (*entity.Organizat
return m.Org, nil
}

func (m *MockGPUCreateStore) GetOrganizations(_ *store.GetOrganizationsOptions) ([]entity.Organization, error) {
if m.Org != nil {
return []entity.Organization{*m.Org}, nil
}
return []entity.Organization{}, nil
}

func (m *MockGPUCreateStore) GetWorkspace(workspaceID string) (*entity.Workspace, error) {
if ws, ok := m.Workspaces[workspaceID]; ok {
return ws, nil
Expand Down
2 changes: 1 addition & 1 deletion pkg/cmd/notebook/notebook.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ func NewCmdNotebook(store NotebookStore, t *terminal.Terminal) *cobra.Command {
hello.TypeItToMeUnskippable27("\nClick here to go to your Jupyter notebook:\n\t 👉" + urlType("http://localhost:8888") + "👈\n\n\n")

// Port forward on 8888
err2 := portforward.RunPortforward(t, store, args[0], "8888:8888", false)
err2 := portforward.RunPortforward(t, store, args[0], "8888:8888", false, "")
if err2 != nil {
return breverrors.WrapAndTrace(err2)
}
Expand Down
9 changes: 6 additions & 3 deletions pkg/cmd/ollama/ollama.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ type OllamaStore interface {
refresh.RefreshStore
util.GetWorkspaceByNameOrIDErrStore
GetActiveOrganizationOrDefault() (*entity.Organization, error)
GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error)
GetCurrentUser() (*entity.User, error)
CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error)
GetWorkspace(workspaceID string) (*entity.Workspace, error)
Expand Down Expand Up @@ -71,6 +72,7 @@ func validateModelType(input string) (bool, error) {
func NewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command {
var model string
var gpu string
var org string

cmd := &cobra.Command{
Use: "ollama",
Expand Down Expand Up @@ -106,7 +108,7 @@ func NewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command
err := runOllamaWorkspace(t, RunOptions{
Model: model,
GPUType: gpu,
}, ollamaStore)
}, ollamaStore, org)
if err != nil {
return nil, breverrors.WrapAndTrace(err)
}
Expand All @@ -127,6 +129,7 @@ func NewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command
}
cmd.Flags().StringVarP(&model, "model", "m", "", "AI/ML model type (e.g., llama2, llama3, mistral7b)")
cmd.Flags().StringVarP(&gpu, "gpu", "g", "g5.xlarge", "GPU instance type. See https://brev.dev/docs/reference/gpu for details")
cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)")
return cmd
}

Expand All @@ -135,13 +138,13 @@ type RunOptions struct {
GPUType string
}

func runOllamaWorkspace(t *terminal.Terminal, opts RunOptions, ollamaStore OllamaStore) error { //nolint:funlen, gocyclo // todo
func runOllamaWorkspace(t *terminal.Terminal, opts RunOptions, ollamaStore OllamaStore, orgFlag string) error { //nolint:funlen, gocyclo // todo
_, err := ollamaStore.GetCurrentUser()
if err != nil {
return breverrors.WrapAndTrace(err)
}

org, err := ollamaStore.GetActiveOrganizationOrDefault()
org, err := util.ResolveOrgFromFlag(ollamaStore, orgFlag)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand Down
19 changes: 15 additions & 4 deletions pkg/cmd/open/open.go
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,7 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto
var host bool
var setDefault string
var editor string
var org string

cmd := &cobra.Command{
Annotations: map[string]string{"access": ""},
Expand Down Expand Up @@ -162,7 +163,7 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto
if len(instanceNames) > 1 {
fmt.Fprintf(os.Stderr, "Opening %s...\n", instanceName)
}
err = runOpenCommand(t, store, instanceName, setupDoneString, directory, host, editorType)
err = runOpenCommand(t, store, instanceName, setupDoneString, directory, host, editorType, org)
if err != nil {
if len(instanceNames) > 1 {
fmt.Fprintf(os.Stderr, "Error opening %s: %v\n", instanceName, err)
Expand All @@ -187,6 +188,12 @@ func NewCmdOpen(t *terminal.Terminal, store OpenStore, noLoginStartStore OpenSto
cmd.Flags().StringVarP(&directory, "dir", "d", "", "directory to open")
cmd.Flags().StringVar(&setDefault, "set-default", "", "set default editor (code, cursor, windsurf, terminal, or tmux)")
cmd.Flags().StringVarP(&editor, "editor", "e", "", "editor to use (code, cursor, windsurf, terminal, or tmux)")
cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)")
errRegComp := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noLoginStartStore, t))
if errRegComp != nil {
breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(errRegComp))
fmt.Print(breverrors.WrapAndTrace(errRegComp))
}

return cmd
}
Expand Down Expand Up @@ -279,11 +286,15 @@ func handleSetDefault(t *terminal.Terminal, editorType string) error {
}

// Fetch workspace info, then open code editor
func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, setupDoneString string, directory string, host bool, editorType string) error { //nolint:funlen,gocyclo // define brev command
func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, setupDoneString string, directory string, host bool, editorType string, orgFlag string) error { //nolint:funlen,gocyclo // define brev command
// todo check if workspace is stopped and start if it if it is stopped
fmt.Println("finding your instance...")
res := refresh.RunRefreshAsync(tstore)
target, err := util.ResolveWorkspaceOrNode(tstore, wsIDOrName)
resolvedOrg, err := util.ResolveOrgFromFlag(tstore, orgFlag)
if err != nil {
return breverrors.WrapAndTrace(err)
}
target, err := util.ResolveWorkspaceOrNodeInOrg(tstore, wsIDOrName, resolvedOrg.ID)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand All @@ -306,7 +317,7 @@ func runOpenCommand(t *terminal.Terminal, tstore OpenStore, wsIDOrName string, s
return breverrors.WrapAndTrace(err)
}

workspace, err = util.GetUserWorkspaceByNameOrIDErr(tstore, wsIDOrName)
workspace, err = util.GetUserWorkspaceByNameOrIDErrInOrg(tstore, wsIDOrName, resolvedOrg.ID)
if err != nil {
return breverrors.WrapAndTrace(err)
}
Expand Down
Loading
Loading