Implement Branch Connection View for Postgres and Vitess

.gitignore

@@ -16,6 +16,7 @@ dist/
 completions/
 vendor/
 .idea
+.worktrees/
 
 # Output of running go build cmd/pscale/main.go
 main

go.mod

@@ -9,7 +9,11 @@ require (
 	github.com/adrg/xdg v0.5.3
 	github.com/benbjohnson/clock v1.3.5
 	github.com/briandowns/spinner v1.23.2
+	github.com/charmbracelet/bubbles v0.21.1-0.20250623103423-23b8fd6302d7
+	github.com/charmbracelet/bubbletea v1.3.10
 	github.com/charmbracelet/huh v1.0.0
+	github.com/charmbracelet/lipgloss v1.1.0
+	github.com/charmbracelet/x/ansi v0.11.2
 	github.com/fatih/color v1.19.0
 	github.com/frankban/quicktest v1.14.6
 	github.com/go-sql-driver/mysql v1.9.3
@@ -24,6 +28,7 @@ require (
 	github.com/mattn/go-isatty v0.0.20
 	github.com/mattn/go-shellwords v1.0.12
 	github.com/mitchellh/go-homedir v1.1.0
+	github.com/muesli/termenv v0.16.0
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
 	github.com/planetscale/planetscale-go v0.168.1
 	github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
@@ -53,11 +58,7 @@ require (
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/catppuccin/go v0.3.0 // indirect
-	github.com/charmbracelet/bubbles v0.21.1-0.20250623103423-23b8fd6302d7 // indirect
-	github.com/charmbracelet/bubbletea v1.3.10 // indirect
 	github.com/charmbracelet/colorprofile v0.3.3 // indirect
-	github.com/charmbracelet/lipgloss v1.1.0 // indirect
-	github.com/charmbracelet/x/ansi v0.11.2 // indirect
 	github.com/charmbracelet/x/cellbuf v0.0.14 // indirect
 	github.com/charmbracelet/x/exp/strings v0.0.0-20251201173703-9f73bfd934ff // indirect
 	github.com/charmbracelet/x/term v0.2.2 // indirect
@@ -93,7 +94,6 @@ require (
 	github.com/mtibben/percent v0.2.1 // indirect
 	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
-	github.com/muesli/termenv v0.16.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.4 // indirect
 	github.com/pires/go-proxyproto v0.8.1 // indirect
 	github.com/planetscale/vitess-types v0.0.0-20250728133330-81b28fd54ee5 // indirect

internal/cmd/branch/branch.go

@@ -35,6 +35,7 @@ func BranchCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.AddCommand(RoutingRulesCmd(ch))
 	cmd.AddCommand(SafeMigrationsCmd(ch))
 	cmd.AddCommand(LintCmd(ch))
+	cmd.AddCommand(ConnectionsCmd(ch))
 	cmd.AddCommand(ProcesslistCmd(ch))
 	cmd.AddCommand(vtctld.VtctldCmd(ch))
 	cmd.AddCommand(InfraCmd(ch))

internal/cmd/branch/connections.go

@@ -0,0 +1,40 @@
+package branch
+
+import (
+	"github.com/planetscale/cli/internal/cmd/branch/connections"
+	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/spf13/cobra"
+)
+
+// ConnectionsCmd manages branch connections across supported database engines.
+func ConnectionsCmd(ch *cmdutil.Helper) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "connections <command>",
+		Short: "Show and kill branch connections",
+		Long: `Show and kill branch connections.
+
+Agent workflow:
+  1. Run: pscale branch connections show <database> <branch> --format json
+  2. Inspect query_id, transaction_id, and connection_id from the selected row.
+  3. Explain the proposed action and wait for user approval before running it.
+  4. Run exactly one action command with the matching ID.
+  5. Run show again to verify the result.
+
+Action semantics:
+  kill <database> <branch> <query-id> --query        Cancels the listed query_id.
+  kill-transaction <database> <branch> <transaction-id>
+                                                     Postgres only. destructive. Terminates the listed transaction_id if it still matches server state.
+  kill <database> <branch> <connection-id>           destructive. Terminates the listed connection_id.
+
+Use --format json when an agent or script needs to inspect query_id,
+transaction_id, and connection_id fields. Human output uses vertical records so
+query text and action IDs are not truncated.`,
+	}
+
+	cmd.AddCommand(ConnectionsShowCmd(ch))
+	cmd.AddCommand(ConnectionsKillCmd(ch))
+	cmd.AddCommand(ConnectionsKillTransactionCmd(ch))
+	cmd.AddCommand(connections.TopCmd(ch))
+
+	return cmd
+}

internal/cmd/branch/connections/actions.go

@@ -0,0 +1,162 @@
+package connections
+
+import (
+	"context"
+	"errors"
+	"strings"
+
+	"github.com/planetscale/cli/internal/cmdutil"
+	live "github.com/planetscale/cli/internal/connections"
+	"github.com/planetscale/cli/internal/printer"
+	ps "github.com/planetscale/planetscale-go/planetscale"
+)
+
+type actionResult struct {
+	Success  bool   `csv:"success" header:"success" json:"success"`
+	Keyspace string `csv:"keyspace" header:"keyspace" json:"keyspace,omitempty"`
+	Shard    string `csv:"shard" header:"shard" json:"shard,omitempty"`
+	Tablet   string `csv:"tablet" header:"tablet" json:"tablet,omitempty"`
+	ID       int64  `csv:"id" header:"id,text" json:"id,omitempty"`
+	Kind     string `csv:"kind" header:"kind" json:"kind,omitempty"`
+}
+
+func (a *actionResult) MarshalCSVValue() interface{} {
+	return []*actionResult{a}
+}
+
+type compactActionResult struct {
+	Success bool   `csv:"success" header:"success" json:"success"`
+	ID      int64  `csv:"id" header:"id,text" json:"id,omitempty"`
+	Kind    string `csv:"kind" header:"kind" json:"kind,omitempty"`
+}
+
+func (a *compactActionResult) MarshalCSVValue() interface{} {
+	return []*compactActionResult{a}
+}
+
+func toActionResult(result live.ActionResult) *actionResult {
+	return &actionResult{
+		Success:  result.Success,
+		Keyspace: result.Keyspace,
+		Shard:    result.Shard,
+		Tablet:   result.Tablet,
+		ID:       result.ID,
+		Kind:     result.Kind,
+	}
+}
+
+func toCompactActionResult(result live.ActionResult) *compactActionResult {
+	return &compactActionResult{
+		Success: result.Success,
+		ID:      result.ID,
+		Kind:    result.Kind,
+	}
+}
+
+// RunCancelQuery cancels the active query identified by a live connection query ID.
+func RunCancelQuery(ctx context.Context, ch *cmdutil.Helper, database, branch, queryID string, target ConnectionTarget) error {
+	return RunCancelQueryForEngine(ctx, ch, database, branch, queryID, ps.DatabaseEngineMySQL, target)
+}
+
+// RunCancelQueryForEngine cancels the active query and prints output for the resolved database engine.
+func RunCancelQueryForEngine(ctx context.Context, ch *cmdutil.Helper, database, branch, queryID string, engine ps.DatabaseEngine, target ConnectionTarget) error {
+	return runAction(ctx, ch, database, branch, "query-id", queryID, target, func(ctx context.Context, client *live.Client, id string) (live.ActionResult, error) {
+		return client.CancelQueryResult(ctx, live.ActionTarget{QueryID: &id})
+	}, engine)
+}
+
+// RunKillTransaction terminates the connection identified by a live connection transaction ID.
+func RunKillTransaction(ctx context.Context, ch *cmdutil.Helper, database, branch, transactionID string, target ConnectionTarget) error {
+	return RunKillTransactionForEngine(ctx, ch, database, branch, transactionID, ps.DatabaseEnginePostgres, target)
+}
+
+// RunKillTransactionForEngine terminates a transaction and prints output for the resolved database engine.
+func RunKillTransactionForEngine(ctx context.Context, ch *cmdutil.Helper, database, branch, transactionID string, engine ps.DatabaseEngine, target ConnectionTarget) error {
+	return runAction(ctx, ch, database, branch, "transaction-id", transactionID, target, func(ctx context.Context, client *live.Client, id string) (live.ActionResult, error) {
+		return client.TerminateTransactionResult(ctx, live.ActionTarget{TransactionID: &id})
+	}, engine)
+}
+
+// RunKillConnection terminates the connection identified by a live connection_id.
+func RunKillConnection(ctx context.Context, ch *cmdutil.Helper, database, branch, connectionID string, target ConnectionTarget) error {
+	return RunKillConnectionForEngine(ctx, ch, database, branch, connectionID, ps.DatabaseEngineMySQL, target)
+}
+
+// RunKillConnectionForEngine terminates a connection and prints output for the resolved database engine.
+func RunKillConnectionForEngine(ctx context.Context, ch *cmdutil.Helper, database, branch, connectionID string, engine ps.DatabaseEngine, target ConnectionTarget) error {
+	return runAction(ctx, ch, database, branch, "connection-id", connectionID, target, func(ctx context.Context, client *live.Client, id string) (live.ActionResult, error) {
+		return client.TerminateConnectionResult(ctx, live.ActionTarget{ConnectionID: &id})
+	}, engine)
+}
+
+func runAction(ctx context.Context, ch *cmdutil.Helper, database, branch, idName, id string, target ConnectionTarget, runAction func(context.Context, *live.Client, string) (live.ActionResult, error), engine ps.DatabaseEngine) error {
+	if err := validateActionID(idName, id); err != nil {
+		return err
+	}
+	id = strings.TrimSpace(id)
+
+	client, err := newConnectionsClient(ch, database, branch, target)
+	if err != nil {
+		return err
+	}
+
+	result, err := runAction(ctx, client, id)
+	if err != nil {
+		return err
+	}
+	return printActionResult(ch, result, engine, idName)
+}
+
+// ValidateConnectionID checks the connection action identifier without making network calls.
+func ValidateConnectionID(id string) error {
+	return validateActionID("connection-id", id)
+}
+
+// ValidateQueryID checks the query action identifier without making network calls.
+func ValidateQueryID(id string) error {
+	return validateActionID("query-id", id)
+}
+
+// ValidateTransactionID checks the transaction action identifier without making network calls.
+func ValidateTransactionID(id string) error {
+	return validateActionID("transaction-id", id)
+}
+
+func validateActionID(idName, id string) error {
+	if strings.TrimSpace(id) == "" {
+		return errors.New(idName + " is required")
+	}
+	return nil
+}
+
+func printActionResult(ch *cmdutil.Helper, result live.ActionResult, engine ps.DatabaseEngine, idName string) error {
+	if ch.Printer.Format() == printer.Human {
+		ch.Printer.Printf("%s.\n", actionResultMessage(result, idName))
+		return nil
+	}
+	if ch.Printer.Format() == printer.JSON {
+		return ch.Printer.PrintResource(toActionResult(result))
+	}
+	if engine == ps.DatabaseEnginePostgres {
+		return ch.Printer.PrintResource(toCompactActionResult(result))
+	}
+	return ch.Printer.PrintResource(toActionResult(result))
+}
+
+func actionResultMessage(result live.ActionResult, idName string) string {
+	var message string
+	switch idName {
+	case "query-id":
+		message = "Cancelled query"
+	case "transaction-id":
+		message = "Killed transaction"
+	case "connection-id":
+		message = "Killed connection"
+	default:
+		message = "Action sent"
+	}
+	if result.Tablet != "" {
+		message += " on " + result.Tablet
+	}
+	return message
+}