Add pscale branch extensions list

AGENTS.md

@@ -334,6 +334,9 @@ pscale maintenance windows <database> <schedule-id> --org <org> --format json
 pscale branch parameters list <database> <branch> --org <org> --format json
 pscale branch parameters list <database> <branch> --org <org> --format json --namespace pgconf
 
+# Extensions available on the cluster image (not CREATE EXTENSION state)
+pscale branch extensions list <database> <branch> --org <org> --format json
+
 # Change parameters (repeat --parameters; keys are namespace.name)
 pscale branch resize <database> <branch> --org <org> --format json --parameters pgconf.max_connections=200
 

internal/cmd/branch/branch.go

@@ -28,6 +28,7 @@ func BranchCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.AddCommand(ResizeCmd(ch))
 	cmd.AddCommand(VtgateCmd(ch))
 	cmd.AddCommand(ParametersCmd(ch))
+	cmd.AddCommand(ExtensionsCmd(ch))
 	cmd.AddCommand(ShowCmd(ch))
 	cmd.AddCommand(UpdateCmd(ch))
 	cmd.AddCommand(SwitchCmd(ch))

internal/cmd/branch/extensions.go

@@ -0,0 +1,110 @@
+package branch
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	ps "github.com/planetscale/cli/internal/planetscale"
+	"github.com/planetscale/cli/internal/printer"
+	"github.com/spf13/cobra"
+)
+
+// ExtensionsCmd lists extensions available on a Postgres branch's cluster image.
+func ExtensionsCmd(ch *cmdutil.Helper) *cobra.Command {
+	long := `List extensions available on a Postgres branch's cluster image.
+
+This is the catalog of extensions the image can load, not the result of
+CREATE EXTENSION. There is no CLI command to enable an extension; preload
+libraries are configured with 'pscale branch resize --parameters'.`
+
+	run := func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+		database, branch := args[0], args[1]
+
+		client, err := ch.Client()
+		if err != nil {
+			return err
+		}
+
+		end := ch.Printer.PrintProgress(fmt.Sprintf("Fetching extensions for branch %s in %s...", printer.BoldBlue(branch), printer.BoldBlue(database)))
+		defer end()
+
+		extensions, err := client.PostgresBranches.ListExtensions(ctx, &ps.ListPostgresExtensionsRequest{
+			Organization: ch.Config.Organization,
+			Database:     database,
+			Branch:       branch,
+		})
+		if err != nil {
+			switch cmdutil.ErrCode(err) {
+			case ps.ErrNotFound:
+				return fmt.Errorf("database %s or branch %s does not exist in organization %s",
+					printer.BoldBlue(database), printer.BoldBlue(branch), printer.BoldBlue(ch.Config.Organization))
+			default:
+				return cmdutil.HandleError(err)
+			}
+		}
+		end()
+
+		if len(extensions) == 0 && ch.Printer.Format() == printer.Human {
+			ch.Printer.Printf("No extensions are listed for %s/%s.\n",
+				printer.BoldBlue(database), printer.BoldBlue(branch))
+			return nil
+		}
+
+		return ch.Printer.PrintResource(toPostgresExtensions(extensions))
+	}
+
+	cmd := &cobra.Command{
+		Use:   "extensions <database> <branch>",
+		Short: "List extensions available on a Postgres branch",
+		Long:  long,
+		Args:  cmdutil.RequiredArgs("database", "branch"),
+		RunE:  run,
+	}
+
+	listCmd := &cobra.Command{
+		Use:     "list <database> <branch>",
+		Short:   "List extensions available on a Postgres branch",
+		Long:    long,
+		Args:    cmdutil.RequiredArgs("database", "branch"),
+		Aliases: []string{"ls"},
+		RunE:    run,
+	}
+	cmd.AddCommand(listCmd)
+
+	return cmd
+}
+
+type postgresExtension struct {
+	Name              string `header:"name" json:"name"`
+	Loader            string `header:"loader" json:"loader"`
+	Available         bool   `header:"available" json:"available"`
+	UnavailableReason string `header:"unavailable,n/a" json:"unavailable_reason"`
+	URL               string `header:"url,n/a" json:"url"`
+
+	orig *ps.PostgresExtension
+}
+
+func toPostgresExtensions(extensions []*ps.PostgresExtension) []*postgresExtension {
+	out := make([]*postgresExtension, 0, len(extensions))
+	for _, ext := range extensions {
+		out = append(out, &postgresExtension{
+			Name:              ext.Name,
+			Loader:            ext.Loader,
+			Available:         ext.Available,
+			UnavailableReason: ext.UnavailableReason,
+			URL:               ext.URL,
+			orig:              ext,
+		})
+	}
+	return out
+}
+
+func (e *postgresExtension) MarshalJSON() ([]byte, error) {
+	return json.MarshalIndent(e.orig, "", "  ")
+}
+
+func (e *postgresExtension) MarshalCSVValue() interface{} {
+	return []*postgresExtension{e}
+}

internal/cmd/branch/extensions_test.go

@@ -0,0 +1,89 @@
+package branch
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	qt "github.com/frankban/quicktest"
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/config"
+	"github.com/planetscale/cli/internal/mock"
+	ps "github.com/planetscale/cli/internal/planetscale"
+	"github.com/planetscale/cli/internal/printer"
+)
+
+func TestBranch_ExtensionsCmd(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "planetscale"
+	db := "postgres-db"
+	branch := "main"
+
+	pgSvc := &mock.PostgresBranchesService{
+		ListExtensionsFn: func(ctx context.Context, req *ps.ListPostgresExtensionsRequest) ([]*ps.PostgresExtension, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			c.Assert(req.Database, qt.Equals, db)
+			c.Assert(req.Branch, qt.Equals, branch)
+			return []*ps.PostgresExtension{
+				{Name: "vector", Loader: "shared_preload_libraries", Available: true, URL: "https://github.com/pgvector/pgvector"},
+				{Name: "pg_stat_statements", Loader: "shared_preload_libraries", Available: false, UnavailableReason: "container_upgrade_required"},
+			}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{PostgresBranches: pgSvc}, nil
+		},
+	}
+
+	cmd := ExtensionsCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(pgSvc.ListExtensionsFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.Contains, "vector")
+	c.Assert(buf.String(), qt.Contains, "pg_stat_statements")
+}
+
+func TestBranch_ExtensionsCmd_ListSubcommand(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	pgSvc := &mock.PostgresBranchesService{
+		ListExtensionsFn: func(ctx context.Context, req *ps.ListPostgresExtensionsRequest) ([]*ps.PostgresExtension, error) {
+			return []*ps.PostgresExtension{
+				{Name: "vector", Available: true},
+			}, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: "planetscale"},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{PostgresBranches: pgSvc}, nil
+		},
+	}
+
+	cmd := ExtensionsCmd(ch)
+	cmd.SetArgs([]string{"list", "postgres-db", "main"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(pgSvc.ListExtensionsFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.Contains, "vector")
+}

internal/mock/branch.go

@@ -204,6 +204,9 @@ type PostgresBranchesService struct {
 
 	ListParametersFn        func(context.Context, *ps.ListPostgresParametersRequest) ([]*ps.PostgresParameter, error)
 	ListParametersFnInvoked bool
+
+	ListExtensionsFn        func(context.Context, *ps.ListPostgresExtensionsRequest) ([]*ps.PostgresExtension, error)
+	ListExtensionsFnInvoked bool
 }
 
 func (p *PostgresBranchesService) Create(ctx context.Context, req *ps.CreatePostgresBranchRequest) (*ps.PostgresBranch, error) {
@@ -265,3 +268,8 @@ func (p *PostgresBranchesService) ListParameters(ctx context.Context, req *ps.Li
 	p.ListParametersFnInvoked = true
 	return p.ListParametersFn(ctx, req)
 }
+
+func (p *PostgresBranchesService) ListExtensions(ctx context.Context, req *ps.ListPostgresExtensionsRequest) ([]*ps.PostgresExtension, error) {
+	p.ListExtensionsFnInvoked = true
+	return p.ListExtensionsFn(ctx, req)
+}

internal/planetscale/postgres_branches.go

@@ -208,6 +208,28 @@ type PostgresParameter struct {
 	UpdatedAt *time.Time `json:"updated_at"`
 }
 
+// ListPostgresExtensionsRequest lists extensions available on a Postgres
+// branch's cluster image.
+type ListPostgresExtensionsRequest struct {
+	Organization string
+	Database     string
+	Branch       string
+}
+
+// PostgresExtension is an extension defined on the branch's cluster image.
+// This is the catalog of what the image can load, not CREATE EXTENSION state.
+type PostgresExtension struct {
+	Type              string               `json:"type"`
+	Name              string               `json:"name"`
+	Description       string               `json:"description"`
+	Internal          bool                 `json:"internal"`
+	Loader            string               `json:"loader"`
+	URL               string               `json:"url"`
+	Available         bool                 `json:"available"`
+	UnavailableReason string               `json:"unavailable_reason"`
+	Parameters        []*PostgresParameter `json:"parameters"`
+}
+
 // PostgresBranchSchemaRequest encapsulates the request to get the schema of a Postgres branch.
 type PostgresBranchSchemaRequest struct {
 	Organization string
@@ -241,6 +263,7 @@ type PostgresBranchesService interface {
 	GetChange(context.Context, *GetPostgresBranchChangeRequest) (*PostgresBranchClusterResizeRequest, error)
 	CancelChanges(context.Context, *CancelPostgresBranchChangesRequest) error
 	ListParameters(context.Context, *ListPostgresParametersRequest) ([]*PostgresParameter, error)
+	ListExtensions(context.Context, *ListPostgresExtensionsRequest) ([]*PostgresExtension, error)
 }
 
 type postgresBranchesService struct {
@@ -474,6 +497,23 @@ func (p *postgresBranchesService) ListParameters(ctx context.Context, listReq *L
 	return parameters, nil
 }
 
+// ListExtensions returns extensions available on the Postgres branch's cluster
+// image. The API returns a bare JSON array.
+func (p *postgresBranchesService) ListExtensions(ctx context.Context, listReq *ListPostgresExtensionsRequest) ([]*PostgresExtension, error) {
+	path := path.Join(postgresBranchAPIPath(listReq.Organization, listReq.Database, listReq.Branch), "extensions")
+	req, err := p.client.newRequest(http.MethodGet, path, nil)
+	if err != nil {
+		return nil, fmt.Errorf("error creating http request: %w", err)
+	}
+
+	extensions := []*PostgresExtension{}
+	if err := p.client.do(ctx, req, &extensions); err != nil {
+		return nil, err
+	}
+
+	return extensions, nil
+}
+
 // Schema returns the schema for the specified Postgres branch.
 func (p *postgresBranchesService) Schema(ctx context.Context, schemaReq *PostgresBranchSchemaRequest) ([]*PostgresBranchSchema, error) {
 	path := path.Join(postgresBranchAPIPath(schemaReq.Organization, schemaReq.Database, schemaReq.Branch), "schema")

internal/planetscale/postgres_branches_test.go

@@ -553,6 +553,35 @@ func TestPostgresBranches_ListParameters(t *testing.T) {
 	c.Assert(parameters[0].Max, qt.Equals, float64(5000))
 }
 
+func TestPostgresBranches_ListExtensions(t *testing.T) {
+	c := qt.New(t)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c.Assert(r.Method, qt.Equals, http.MethodGet)
+		c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/postgres-test-db/branches/postgres-test-branch/extensions")
+		w.WriteHeader(200)
+		out := `[{"type":"PostgresClusterExtension","name":"vector","description":"<p>vector</p>","internal":false,"loader":"shared_preload_libraries","url":"https://github.com/pgvector/pgvector","available":true,"unavailable_reason":"","parameters":[]}]`
+		_, err := w.Write([]byte(out))
+		c.Assert(err, qt.IsNil)
+	}))
+
+	client, err := NewClient(WithBaseURL(ts.URL))
+	c.Assert(err, qt.IsNil)
+
+	extensions, err := client.PostgresBranches.ListExtensions(context.Background(), &ListPostgresExtensionsRequest{
+		Organization: "my-org",
+		Database:     "postgres-test-db",
+		Branch:       testPostgresBranch,
+	})
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(extensions, qt.HasLen, 1)
+	c.Assert(extensions[0].Name, qt.Equals, "vector")
+	c.Assert(extensions[0].Loader, qt.Equals, "shared_preload_libraries")
+	c.Assert(extensions[0].Available, qt.IsTrue)
+	c.Assert(extensions[0].URL, qt.Equals, "https://github.com/pgvector/pgvector")
+}
+
 func TestPostgresBranches_ListParametersWithFilters(t *testing.T) {
 	c := qt.New(t)