From ca919a28685e0e16a2df1c55276fd58022afe2e4 Mon Sep 17 00:00:00 2001
From: eileencodes <eileencodes@gmail.com>
Date: Mon, 8 Jun 2026 16:51:14 -0400
Subject: [PATCH] Raise correct error when a resource doesn't exist

I was doing some testing with the CLI and typed the organization name
incorrectly. Instead of giving me an error the org didn't exist, it told
me the branch didn't exist which was confusing. Change the error
handling to determine what is 404ing to give the user an accurate
representation of what went wrong.
---
 internal/cmd/branch/create.go      | 12 +++-
 internal/cmd/branch/create_test.go | 93 ++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 2 deletions(-)

diff --git a/internal/cmd/branch/create.go b/internal/cmd/branch/create.go
index d650a84a..baf7edc3 100644
--- a/internal/cmd/branch/create.go
+++ b/internal/cmd/branch/create.go
@@ -86,10 +86,18 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 			if err != nil {
 				switch cmdutil.ErrCode(err) {
 				case ps.ErrNotFound:
+					if _, orgErr := client.Organizations.Get(ctx, &ps.GetOrganizationRequest{
+						Organization: ch.Config.Organization,
+					}); cmdutil.ErrCode(orgErr) == ps.ErrNotFound {
+						return cmdutil.HandleNotFoundWithServiceTokenCheck(
+							ctx, cmd, ch.Config, ch.Client, err, "read_branch",
+							"organization %s does not exist",
+							printer.BoldBlue(ch.Config.Organization))
+					}
 					return cmdutil.HandleNotFoundWithServiceTokenCheck(
 						ctx, cmd, ch.Config, ch.Client, err, "read_branch",
-						"branch %s does not exist in database %s (organization: %s)",
-						printer.BoldBlue(branch), printer.BoldBlue(source), printer.BoldBlue(ch.Config.Organization))
+						"source database %s does not exist in organization %s",
+						printer.BoldBlue(source), printer.BoldBlue(ch.Config.Organization))
 				default:
 					return cmdutil.HandleError(err)
 				}
diff --git a/internal/cmd/branch/create_test.go b/internal/cmd/branch/create_test.go
index b41b7031..5b5012fb 100644
--- a/internal/cmd/branch/create_test.go
+++ b/internal/cmd/branch/create_test.go
@@ -529,3 +529,96 @@ func TestBranch_CreateCmd_GenuineNotFound(t *testing.T) {
 	c.Assert(branchSvc.CreateFnInvoked, qt.IsTrue)
 	c.Assert(orgSvc.ListFnInvoked, qt.IsTrue)
 }
+
+func TestBranch_CreateCmd_OrgNotFound(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "nonexistent-org"
+	db := "planetscale"
+	branch := "development"
+
+	orgSvc := &mock.OrganizationsService{
+		GetFn: func(ctx context.Context, req *ps.GetOrganizationRequest) (*ps.Organization, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	dbSvc := &mock.DatabaseService{
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Organizations: orgSvc,
+				Databases:     dbSvc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "organization "+org+" does not exist")
+	c.Assert(err.Error(), qt.Not(qt.Contains), branch)
+	c.Assert(orgSvc.GetFnInvoked, qt.IsTrue)
+}
+
+func TestBranch_CreateCmd_SourceDatabaseNotFound(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "planetscale"
+	db := "nonexistent-db"
+	branch := "development"
+
+	orgSvc := &mock.OrganizationsService{
+		GetFn: func(ctx context.Context, req *ps.GetOrganizationRequest) (*ps.Organization, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			return &ps.Organization{Name: org}, nil
+		},
+	}
+
+	dbSvc := &mock.DatabaseService{
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
+			c.Assert(req.Database, qt.Equals, db)
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Organizations: orgSvc,
+				Databases:     dbSvc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "source database "+db+" does not exist in organization "+org)
+	c.Assert(err.Error(), qt.Not(qt.Contains), branch)
+	c.Assert(orgSvc.GetFnInvoked, qt.IsTrue)
+}