Skip to content

Commit c32ef65

Browse files
eileencodeseileencodes
committed
[Neki] Support reading live router connections in CLI
Adds CLI support for read-only router connections for Neki. Killing connections are not yet supported, nor are backend postgres connections. This functionality is currently behind a feature flag and may not yet be available on your Neki database. Example error if this is the case: ``` Error: live connections are not enabled for branch <branch> in database <database> (organization: <org-name>) ``` Usage examples: ``` $ pscale branch connections show <database> <branch> captured_at: 2026-10-19T19:02:05Z *************************** 1. row *************************** pid: 53 router: default zone: us-east-1a state: idle duration: 1m22.022s user: <redacted> database: postgres application: psql client_addr: <redacted> backend_start: 2026-10-19T18:56:33Z shard_backends: 0 query: <redacted> ``` If there are no queries running on a shard backend, the connections will return `shard_backends: 0`, otherwise it will return information about the shard backend, including the query running. ``` $ pscale branch connections show <database> <branch> captured_at: 2026-10-19T19:07:47Z *************************** 1. row *************************** pid: 35 router: default zone: us-east-1b state: active duration: 6.811s user: <redacted> database: postgres application: psql client_addr: <redacted> backend_start: 2026-10-19T19:07:38Z shard_backends: 1 shard: global shard_uid: <redacted> sidecar: <redacted> backend_pid: 328942 xact_start: 2026-10-19T19:07:41Z query_start: 2026-10-19T19:07:41Z backend_query: <redated> query: <redacted> ```
1 parent 89498e2 commit c32ef65

24 files changed

Lines changed: 1064 additions & 72 deletions

‎internal/cmd/branch/connections.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ import (
1010
func ConnectionsCmd(ch *cmdutil.Helper) *cobra.Command {
1111
cmd := &cobra.Command{
1212
Use: "connections <command>",
13-
Short: "Show and kill MySQL or Postgres branch connections",
14-
Long: `Show and kill MySQL or Postgres branch connections. Neki is not supported.
13+
Short: "Show and kill branch connections",
14+
Long: `Show and kill MySQL or Postgres branch connections, and show Neki branch connections.
1515
1616
MySQL (Vitess):
1717
show/kill/top with --keyspace and --shard.
@@ -20,6 +20,9 @@ Postgres:
2020
show/top with --instance or --role.
2121
kill-transaction terminates a transaction_id from show.
2222
23+
Neki:
24+
show/top list client sessions across the branch's routers. Read-only.
25+
2326
Agent workflow:
2427
1. Run: pscale branch connections show <database> <branch> --format json
2528
2. Inspect query_id, transaction_id, and connection_id from the selected row.

‎internal/cmd/branch/connections/client.go‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@ package connections
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
7+
"sync"
68

79
"github.com/planetscale/cli/internal/cmdutil"
810
live "github.com/planetscale/cli/internal/connections"
@@ -39,6 +41,65 @@ func databaseNotFoundError(ch *cmdutil.Helper, database string) error {
3941
printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization))
4042
}
4143

44+
// notFoundExplainer replaces the connections endpoint's ambiguous 404 with
45+
// whether the branch is missing or just lacks live connections. The answer is
46+
// cached so a refreshing TUI does not repeat the lookups every interval.
47+
type notFoundExplainer struct {
48+
ch *cmdutil.Helper
49+
database string
50+
branch string
51+
52+
once sync.Once
53+
explained error
54+
}
55+
56+
func newNotFoundExplainer(ch *cmdutil.Helper, database, branch string) *notFoundExplainer {
57+
return &notFoundExplainer{ch: ch, database: database, branch: branch}
58+
}
59+
60+
func (e *notFoundExplainer) explain(ctx context.Context, err error) error {
61+
var notFound *live.NotFoundError
62+
if e == nil || !errors.As(err, &notFound) {
63+
return err
64+
}
65+
e.once.Do(func() {
66+
e.explained = e.resolve(ctx, err)
67+
})
68+
return e.explained
69+
}
70+
71+
func (e *notFoundExplainer) resolve(ctx context.Context, err error) error {
72+
client, clientErr := e.ch.Client()
73+
if clientErr != nil {
74+
return err
75+
}
76+
org := e.ch.Config.Organization
77+
db, dbErr := client.Databases.Get(ctx, &ps.GetDatabaseRequest{Organization: org, Database: e.database})
78+
if cmdutil.ErrCode(dbErr) == ps.ErrNotFound {
79+
return databaseNotFoundError(e.ch, e.database)
80+
}
81+
if dbErr != nil || db == nil {
82+
return err
83+
}
84+
85+
var branchErr error
86+
if db.Kind == ps.DatabaseEngineMySQL {
87+
_, branchErr = client.DatabaseBranches.Get(ctx, &ps.GetDatabaseBranchRequest{Organization: org, Database: e.database, Branch: e.branch})
88+
} else {
89+
_, branchErr = client.PostgresBranches.Get(ctx, &ps.GetPostgresBranchRequest{Organization: org, Database: e.database, Branch: e.branch})
90+
}
91+
switch {
92+
case cmdutil.ErrCode(branchErr) == ps.ErrNotFound:
93+
return fmt.Errorf("branch %s does not exist in database %s (organization: %s)",
94+
printer.BoldBlue(e.branch), printer.BoldBlue(e.database), printer.BoldBlue(org))
95+
case branchErr != nil:
96+
return err
97+
default:
98+
return fmt.Errorf("live connections are not enabled for branch %s in database %s (organization: %s)",
99+
printer.BoldBlue(e.branch), printer.BoldBlue(e.database), printer.BoldBlue(org))
100+
}
101+
}
102+
42103
func newConnectionsClient(ch *cmdutil.Helper, database, branch string, target ConnectionTarget) (*live.Client, error) {
43104
return live.NewClient(live.ClientConfig{
44105
BaseURL: ch.Config.BaseURL,

‎internal/cmd/branch/connections/client_test.go‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package connections
2+
3+
import (
4+
"context"
5+
"errors"
6+
"testing"
7+
8+
qt "github.com/frankban/quicktest"
9+
"github.com/planetscale/cli/internal/cmdutil"
10+
"github.com/planetscale/cli/internal/config"
11+
live "github.com/planetscale/cli/internal/connections"
12+
"github.com/planetscale/cli/internal/mock"
13+
ps "github.com/planetscale/cli/internal/planetscale"
14+
)
15+
16+
func TestNotFoundExplainerLooksUpBranchOnce(t *testing.T) {
17+
c := qt.New(t)
18+
lookups := 0
19+
branches := &mock.DatabaseBranchesService{
20+
GetFn: func(context.Context, *ps.GetDatabaseBranchRequest) (*ps.DatabaseBranch, error) {
21+
lookups++
22+
return &ps.DatabaseBranch{Name: "main"}, nil
23+
},
24+
}
25+
ch := &cmdutil.Helper{
26+
Config: &config.Config{Organization: "acme"},
27+
Client: func() (*ps.Client, error) {
28+
return &ps.Client{
29+
Databases: &mock.DatabaseService{
30+
GetFn: func(_ context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
31+
return &ps.Database{Name: req.Database, Kind: ps.DatabaseEngineMySQL}, nil
32+
},
33+
},
34+
DatabaseBranches: branches,
35+
}, nil
36+
},
37+
}
38+
explainer := newNotFoundExplainer(ch, "shop", "main")
39+
notFound := &live.NotFoundError{Organization: "acme", Database: "shop", Branch: "main"}
40+
41+
first := explainer.explain(context.Background(), notFound)
42+
second := explainer.explain(context.Background(), notFound)
43+
44+
c.Assert(first, qt.ErrorMatches, `live connections are not enabled for branch main in database shop \(organization: acme\)`)
45+
c.Assert(second, qt.Equals, first)
46+
c.Assert(lookups, qt.Equals, 1)
47+
}
48+
49+
func TestNotFoundExplainerPassesOtherErrorsThrough(t *testing.T) {
50+
c := qt.New(t)
51+
other := errors.New("list connections: rate limited, please retry in a moment")
52+
53+
got := newNotFoundExplainer(&cmdutil.Helper{}, "shop", "main").explain(context.Background(), other)
54+
55+
c.Assert(got, qt.Equals, other)
56+
}

‎internal/cmd/branch/connections/engine_flags.go‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,26 +6,43 @@ import (
66
ps "github.com/planetscale/cli/internal/planetscale"
77
)
88

9-
var errNekiUnsupported = errors.New("connections is not supported for Neki databases")
9+
var (
10+
errVitessTargetOnly = errors.New("--keyspace/--shard are only supported for Vitess databases")
11+
errPostgresFilterOnly = errors.New("--instance/--role are only supported for Postgres databases")
12+
errNekiReadOnly = errors.New("connections are read-only for Neki databases")
13+
)
1014

1115
// ValidateEngineFlags rejects flags that only apply to another database engine.
1216
func ValidateEngineFlags(engine ps.DatabaseEngine, filter ConnectionFilter, target ConnectionTarget) error {
1317
return validateEngineFlags(engine, filter.connectionFilter(), target)
1418
}
1519

1620
func validateEngineFlags(engine ps.DatabaseEngine, filter connectionFilter, target ConnectionTarget) error {
17-
if engine == ps.DatabaseEngineNeki {
18-
return errNekiUnsupported
19-
}
21+
vitessTarget := target.Keyspace != "" || target.Shard != ""
2022
switch engine {
2123
case ps.DatabaseEnginePostgres:
22-
if target.Keyspace != "" || target.Shard != "" {
23-
return errors.New("--keyspace/--shard are only supported for Vitess databases")
24+
if vitessTarget {
25+
return errVitessTargetOnly
2426
}
2527
case ps.DatabaseEngineMySQL:
2628
if filter.active() {
27-
return errors.New("--instance/--role are only supported for Postgres databases")
29+
return errPostgresFilterOnly
30+
}
31+
case ps.DatabaseEngineNeki:
32+
if vitessTarget {
33+
return errVitessTargetOnly
34+
}
35+
if filter.active() {
36+
return errPostgresFilterOnly
2837
}
2938
}
3039
return nil
3140
}
41+
42+
// ValidateActionEngine rejects kill commands for engines whose connections are read-only.
43+
func ValidateActionEngine(engine ps.DatabaseEngine) error {
44+
if engine == ps.DatabaseEngineNeki {
45+
return errNekiReadOnly
46+
}
47+
return nil
48+
}

‎internal/cmd/branch/connections/engine_flags_test.go‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,38 @@ import (
77
ps "github.com/planetscale/cli/internal/planetscale"
88
)
99

10-
func TestValidateEngineFlagsRejectsNeki(t *testing.T) {
11-
c := qt.New(t)
10+
func TestValidateEngineFlagsForNeki(t *testing.T) {
11+
tests := []struct {
12+
name string
13+
filter ConnectionFilter
14+
target ConnectionTarget
15+
wantErr string
16+
}{
17+
{name: "no flags"},
18+
{name: "vitess target", target: ConnectionTarget{Keyspace: "commerce"}, wantErr: "--keyspace/--shard are only supported for Vitess databases"},
19+
{name: "instance filter", filter: ConnectionFilter{Instance: "router-a"}, wantErr: "--instance/--role are only supported for Postgres databases"},
20+
{name: "role filter", filter: ConnectionFilter{Role: "primary"}, wantErr: "--instance/--role are only supported for Postgres databases"},
21+
}
22+
23+
for _, tt := range tests {
24+
t.Run(tt.name, func(t *testing.T) {
25+
c := qt.New(t)
26+
27+
err := ValidateEngineFlags(ps.DatabaseEngineNeki, tt.filter, tt.target)
1228

13-
err := ValidateEngineFlags(ps.DatabaseEngineNeki, ConnectionFilter{}, ConnectionTarget{})
29+
if tt.wantErr == "" {
30+
c.Assert(err, qt.IsNil)
31+
return
32+
}
33+
c.Assert(err, qt.ErrorMatches, tt.wantErr)
34+
})
35+
}
36+
}
37+
38+
func TestValidateActionEngine(t *testing.T) {
39+
c := qt.New(t)
1440

15-
c.Assert(err, qt.ErrorMatches, "connections is not supported for Neki databases")
41+
c.Assert(ValidateActionEngine(ps.DatabaseEngineNeki), qt.ErrorMatches, "connections are read-only for Neki databases")
42+
c.Assert(ValidateActionEngine(ps.DatabaseEnginePostgres), qt.IsNil)
43+
c.Assert(ValidateActionEngine(ps.DatabaseEngineMySQL), qt.IsNil)
1644
}

‎internal/cmd/branch/connections/replay_test.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,46 @@ func TestTopCmdReplayUsesVitessViewForMySQLTrace(t *testing.T) {
214214
c.Assert(view, qt.Not(qt.Contains), "WAIT")
215215
}
216216

217+
func TestTopCmdReplayUsesNekiViewForNekiTrace(t *testing.T) {
218+
c := qt.New(t)
219+
restoreTTY := setPrinterTTY(t, true)
220+
defer restoreTTY()
221+
222+
var capturedModel tea.Model
223+
restoreProgram := setRunTeaProgram(t, func(model tea.Model, options ...tea.ProgramOption) error {
224+
capturedModel = model
225+
return nil
226+
})
227+
defer restoreProgram()
228+
229+
path := filepath.Join(t.TempDir(), "neki.jsonl")
230+
file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o600)
231+
c.Assert(err, qt.IsNil)
232+
list := live.NewConnectionList(time.Date(2026, 10, 16, 12, 0, 5, 0, time.UTC), []live.Connection{{
233+
PID: 7,
234+
Instance: "router-a",
235+
RouterName: "router-a",
236+
State: "active",
237+
QueryText: "SELECT 1",
238+
}}, live.SortByTransactionStart)
239+
list.DatabaseKind = live.DatabaseKindNeki
240+
list.Warnings = []string{"router router-b unreachable"}
241+
c.Assert(history.NewCaptureWriter(file).Write(history.NewCapture(list)), qt.IsNil)
242+
c.Assert(file.Close(), qt.IsNil)
243+
244+
cmd := testTopCmd(&cmdutil.Helper{Config: &config.Config{}})
245+
cmd.SetArgs([]string{"--replay", path})
246+
247+
c.Assert(cmd.Execute(), qt.IsNil)
248+
sized, _ := capturedModel.(tui.Model).Update(tea.WindowSizeMsg{Width: 160, Height: 24})
249+
view := sized.(tui.Model).View()
250+
251+
c.Assert(view, qt.Contains, "ROUTER")
252+
c.Assert(view, qt.Contains, "router-a")
253+
c.Assert(view, qt.Contains, "incomplete results: router router-b unreachable")
254+
c.Assert(view, qt.Not(qt.Contains), "BLOCK")
255+
}
256+
217257
func TestReplayHeaderShowsTraceTarget(t *testing.T) {
218258
c := qt.New(t)
219259
restoreTTY := setPrinterTTY(t, true)

‎internal/cmd/branch/connections/show.go‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ func runList(ctx context.Context, ch *cmdutil.Helper, database, branch string, f
4747
return err
4848
}
4949

50-
list, err := filteredLister{client: client, filter: filter}.List(ctx, live.SortByTransactionStart)
50+
list, err := filteredLister{client: client, filter: filter, notFound: newNotFoundExplainer(ch, database, branch)}.List(ctx, live.SortByTransactionStart)
5151
if err != nil {
5252
return live.UserFacingError(err, "view")
5353
}
@@ -76,6 +76,7 @@ type printableList struct {
7676
Topology *ListTopology `json:"topology,omitempty"`
7777
Instances []printableInstance `json:"instances"`
7878
Connections []printableConnection `json:"connections"`
79+
Warnings []string `json:"warnings,omitempty"`
7980
}
8081

8182
type printableInstance struct {
@@ -101,6 +102,13 @@ type printableConnection struct {
101102
QueryID *string `csv:"query_id" header:"query_id" json:"query_id"`
102103
TransactionID *string `csv:"transaction_id" header:"transaction_id" json:"transaction_id"`
103104
ConnectionID *string `csv:"connection_id" header:"connection_id" json:"connection_id"`
105+
106+
RouterName string `csv:"-" json:"router_name,omitempty"`
107+
RouterCell string `csv:"-" json:"router_cell,omitempty"`
108+
RouterUID string `csv:"-" json:"router_uid,omitempty"`
109+
AvailabilityZone string `csv:"-" json:"availability_zone,omitempty"`
110+
BackendStart *time.Time `csv:"-" json:"backend_start,omitempty"`
111+
SidecarBackends []live.SidecarBackend `csv:"-" json:"sidecar_backends,omitempty"`
104112
}
105113

106114
func toPrintableList(list live.ConnectionList, topology ListTopology) printableList {
@@ -109,6 +117,7 @@ func toPrintableList(list live.ConnectionList, topology ListTopology) printableL
109117
CapturedAt: list.CapturedAt,
110118
Instances: toPrintableInstances(list.Instances),
111119
Connections: toPrintableConnections(list),
120+
Warnings: list.Warnings,
112121
}
113122
if !topology.isEmpty() {
114123
out.Topology = &topology
@@ -150,7 +159,7 @@ func (p printableList) hasDatabaseName() bool {
150159
func toPrintableConnections(list live.ConnectionList) []printableConnection {
151160
connections := make([]printableConnection, 0, len(list.Connections))
152161
for _, conn := range list.Connections {
153-
connections = append(connections, printableConnection{
162+
printable := printableConnection{
154163
PID: conn.PID,
155164
Instance: conn.Instance,
156165
InstanceRole: conn.InstanceRole,
@@ -167,7 +176,16 @@ func toPrintableConnections(list live.ConnectionList) []printableConnection {
167176
QueryID: conn.QueryID,
168177
TransactionID: conn.TransactionID,
169178
ConnectionID: conn.ConnectionID,
170-
})
179+
}
180+
if list.DatabaseKind == live.DatabaseKindNeki {
181+
printable.RouterName = conn.RouterName
182+
printable.RouterCell = conn.RouterCell
183+
printable.RouterUID = conn.RouterUID
184+
printable.AvailabilityZone = conn.AvailabilityZone
185+
printable.BackendStart = conn.BackendStart
186+
printable.SidecarBackends = conn.SidecarBackends
187+
}
188+
connections = append(connections, printable)
171189
}
172190
return connections
173191
}
@@ -185,6 +203,7 @@ func printableBlockedBy(blockedBy []int) []int {
185203
func printHumanConnectionList(out io.Writer, list live.ConnectionList, topology ListTopology) {
186204
fmt.Fprintf(out, "captured_at: %s\n", list.CapturedAt.Format(time.RFC3339))
187205
vitess := list.DatabaseKind == live.DatabaseKindMySQL || !topology.isEmpty()
206+
neki := list.DatabaseKind == live.DatabaseKindNeki
188207
if !topology.isEmpty() {
189208
fmt.Fprintln(out, "topology:")
190209
fmt.Fprintf(out, " keyspace: %s\n", topology.Keyspace)
@@ -194,6 +213,9 @@ func printHumanConnectionList(out io.Writer, list live.ConnectionList, topology
194213
if warning := unreachableInstanceWarning(list.Instances); warning != "" {
195214
fmt.Fprintln(out, warning)
196215
}
216+
for _, warning := range list.Warnings {
217+
fmt.Fprintln(out, "warning: incomplete results: "+live.SanitizeDisplayText(warning))
218+
}
197219

198220
if len(list.Connections) == 0 {
199221
fmt.Fprintln(out, "No live connections found.")
@@ -207,13 +229,16 @@ func printHumanConnectionList(out io.Writer, list live.ConnectionList, topology
207229

208230
fmt.Fprintf(out, "*************************** %d. row ***************************\n", i+1)
209231
fields := conn.HumanFields()
210-
if vitess {
232+
switch {
233+
case vitess:
211234
fields = vitessHumanFields(conn)
235+
case neki:
236+
fields = conn.NekiHumanFields()
212237
}
213238
for _, field := range fields {
214239
writeHumanField(out, field[0], field[1])
215240
}
216-
if conn.DatabaseName != "" {
241+
if conn.DatabaseName != "" && !neki {
217242
writeHumanField(out, "database", conn.DatabaseName)
218243
}
219244
fmt.Fprintln(out, "query:")

0 commit comments

Comments
 (0)