diff --git a/AGENTS.md b/AGENTS.md index bbff3ae9..f91e5ba2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -162,7 +162,7 @@ pscale sql --org --format json --keyspace - | `--role` | `reader` (default), `writer`, `readwriter`, `admin` — same names as `pscale shell` | | `--replica` | Route reads to replicas | | `--dbname` | PostgreSQL database name (default `postgres`) | -| `--keyspace` | MySQL keyspace (default `@primary`) | +| `--keyspace` | MySQL keyspace (default `@primary`); may include a shard and tablet type: `mykeyspace/-80`, `mykeyspace/-80@replica` | | `--force` | Allow destructive SQL after explicit user approval | **`--role` by engine** (same as `pscale shell`): @@ -197,6 +197,33 @@ Error: one JSON object on stdout with `status: "error"`, `error`, `issues`, and Destructive SQL without `--force`: `status: "action_required"`, `query_kind: "destructive"`, `issues`, and `next_steps` (includes `--force` retry command). +## Diagnostics: insights + inspect + +Two complementary read-only surfaces. When diagnosing database health or performance, **check both** — they see different things. + +**`pscale insights`** — server-side analysis computed from production traffic (works even when you can't or don't want to connect to the database): + +```bash +pscale insights queries --org --format json --sort totalTime # top queries; sorts: totalTime, count, p99Latency, rowsRead, rowsReadPerReturned, errorCount, ... +pscale insights errors --org --format json # failing queries with error messages +pscale insights anomalies --org --format json # detected resource anomalies (CPU, memory, IOPS, rows) +pscale insights recommendations --org --format json # schema recommendations with ready-to-apply DDL +``` + +**`pscale inspect`** — live, point-in-time checks run over a direct connection (same credentials model as `pscale sql`, always read-only): + +```bash +pscale inspect all --org --format json # every applicable check, one report +pscale inspect --org --format json +``` + +Checks: `table-sizes`, `index-sizes`, `unused-indexes`, `redundant-indexes`, `seq-scans`, `long-running-queries`, `locks`, `outliers`, `calls`, `bloat`, `vacuum-stats`, `replication-slots`, `subscriptions`. Checks adapt per engine; ones that don't apply explain the alternative. JSON results include `next_steps` pointing at the matching `insights` command — follow them. + +Caveats: +- Statistics are since last server restart and per-connection-target: on sharded Vitess databases they reflect a single shard's MySQL instance. Use `--keyspace` to pick a keyspace, or pin an exact shard with `--keyspace 'mykeyspace/-80'` (enumerate with `pscale sql --org --format json --query "SHOW VITESS_SHARDS"`; rows are `keyspace/shard`). Databases can have hundreds of shards — inspect one shard at a time rather than fanning out. On PostgreSQL, stats are scoped to one database (use `--dbname`; if CONNECT is denied, retry with `--role admin`). +- `outliers`/`calls` need `pg_stat_statements` on PostgreSQL; if missing, use `pscale insights queries` instead (no extension needed). +- Rule of thumb: start with `insights` (traffic-aware, historical), use `inspect` for live state (locks, in-flight queries) and physical layout (sizes, bloat, index usage). + ## Postgres branch changes (size, replicas, parameters) `pscale branch resize` queues a single asynchronous **change request** for a Postgres branch covering cluster size, replica count, and configuration parameters in any combination. Track it with `resize status`; cancel it with `resize cancel` while queued. diff --git a/internal/cmd/insights/anomalies.go b/internal/cmd/insights/anomalies.go new file mode 100644 index 00000000..8d2e2007 --- /dev/null +++ b/internal/cmd/insights/anomalies.go @@ -0,0 +1,80 @@ +package insights + +import ( + "fmt" + + "github.com/spf13/cobra" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" +) + +// AnomalyRow is an anomaly formatted for table output. +type AnomalyRow struct { + ID string `header:"id" json:"id"` + Active bool `header:"active" json:"active"` + PeriodStart string `header:"started" json:"period_start"` + PeriodEnd string `header:"ended" json:"period_end"` + MinutesInViolation int64 `header:"minutes in violation" json:"minutes_in_violation"` +} + +// AnomaliesCmd lists detected resource anomalies for a branch. +func AnomaliesCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "anomalies ", + Short: "List detected resource anomalies (CPU, memory, IOPS, rows read/written)", + Args: cmdutil.RequiredArgs("database", "branch"), + RunE: 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 anomalies for %s in %s...", + printer.BoldBlue(branch), printer.BoldBlue(database))) + defer end() + + anomalies, err := client.QueryInsights.ListAnomalies(ctx, &ps.ListAnomaliesRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + }) + if err != nil { + return notFoundError(ch, err, database, branch) + } + end() + + if len(anomalies) == 0 && ch.Printer.Format() == printer.Human { + ch.Printer.Printf("No anomalies detected for %s in %s.\n", + printer.BoldBlue(branch), printer.BoldBlue(database)) + return nil + } + + if ch.Printer.Format() == printer.JSON { + return ch.Printer.PrintJSON(anomalies) + } + + rows := make([]*AnomalyRow, 0, len(anomalies)) + for _, a := range anomalies { + end := "" + if !a.PeriodEnd.IsZero() { + end = a.PeriodEnd.Format("2006-01-02 15:04") + } + rows = append(rows, &AnomalyRow{ + ID: a.ID, + Active: a.Active, + PeriodStart: a.PeriodStart.Format("2006-01-02 15:04"), + PeriodEnd: end, + MinutesInViolation: a.MinutesInViolation, + }) + } + return ch.Printer.PrintResource(rows) + }, + } + + return cmd +} diff --git a/internal/cmd/insights/errors.go b/internal/cmd/insights/errors.go new file mode 100644 index 00000000..2d4da27a --- /dev/null +++ b/internal/cmd/insights/errors.go @@ -0,0 +1,82 @@ +package insights + +import ( + "fmt" + + "github.com/spf13/cobra" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" +) + +// ErrorRow is a query error row formatted for table output. +type ErrorRow struct { + Count int64 `header:"count" json:"error_count"` + LastSeen string `header:"last seen" json:"started_at"` + Message string `header:"message" json:"error_message"` + ID string `header:"id" json:"id"` +} + +// ErrorsCmd lists aggregated query errors for a branch. +func ErrorsCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + limit int + period string + } + + cmd := &cobra.Command{ + Use: "errors ", + Short: "List queries that are failing with errors", + Args: cmdutil.RequiredArgs("database", "branch"), + RunE: 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 query errors for %s in %s...", + printer.BoldBlue(branch), printer.BoldBlue(database))) + defer end() + + errs, err := client.QueryInsights.ListErrors(ctx, &ps.ListQueryInsightsErrorsRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + }, ps.WithPerPage(flags.limit), ps.WithPeriod(flags.period)) + if err != nil { + return notFoundError(ch, err, database, branch) + } + end() + + if len(errs) == 0 && ch.Printer.Format() == printer.Human { + ch.Printer.Printf("No query errors recorded for %s in %s.\n", + printer.BoldBlue(branch), printer.BoldBlue(database)) + return nil + } + + if ch.Printer.Format() == printer.JSON { + return ch.Printer.PrintJSON(errs) + } + + rows := make([]*ErrorRow, 0, len(errs)) + for _, e := range errs { + rows = append(rows, &ErrorRow{ + Count: e.ErrorCount, + LastSeen: e.StartedAt.Format("2006-01-02 15:04"), + Message: truncate(e.ErrorMessage, 100), + ID: e.ID, + }) + } + return ch.Printer.PrintResource(rows) + }, + } + + cmd.Flags().IntVar(&flags.limit, "limit", 15, "Number of errors to return") + cmd.Flags().StringVar(&flags.period, "period", "", "Time period to aggregate over (e.g. 1h, 24h)") + + return cmd +} diff --git a/internal/cmd/insights/insights.go b/internal/cmd/insights/insights.go new file mode 100644 index 00000000..e303cb6e --- /dev/null +++ b/internal/cmd/insights/insights.go @@ -0,0 +1,62 @@ +package insights + +import ( + "fmt" + + "github.com/spf13/cobra" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" +) + +// InsightsCmd surfaces server-side query insights for a database: aggregated +// query statistics, query errors, anomalies, and schema recommendations. +func InsightsCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "insights ", + Short: "Query performance insights and schema recommendations for a database", + Long: `Surface PlanetScale's server-side analysis of a database: aggregated query +statistics (latency percentiles, rows read, errors), detected anomalies, and +schema recommendations, all computed from production traffic. + +For live, connection-level diagnostics (table sizes, locks, running queries), +see pscale inspect.`, + PersistentPreRunE: cmdutil.CheckAuthentication(ch.Config), + } + + cmd.PersistentFlags().StringVar(&ch.Config.Organization, "org", ch.Config.Organization, + "The organization for the current user") + cmd.MarkPersistentFlagRequired("org") // nolint:errcheck + + cmd.AddCommand(QueriesCmd(ch)) + cmd.AddCommand(ErrorsCmd(ch)) + cmd.AddCommand(AnomaliesCmd(ch)) + cmd.AddCommand(RecommendationsCmd(ch)) + + return cmd +} + +// notFoundError maps a not-found API error to a message explaining both +// causes: a missing database/branch, or insights not being enabled. +func notFoundError(ch *cmdutil.Helper, err error, database, branch string) error { + switch cmdutil.ErrCode(err) { + case ps.ErrNotFound: + if branch != "" { + return fmt.Errorf("branch %s does not exist in database %s (organization: %s) or query insights is not enabled for the database", + printer.BoldBlue(branch), printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization)) + } + return fmt.Errorf("database %s does not exist in organization %s", + printer.BoldBlue(database), printer.BoldBlue(ch.Config.Organization)) + default: + return cmdutil.HandleError(err) + } +} + +// truncate shortens s for table output. +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n-1] + "…" +} diff --git a/internal/cmd/insights/queries.go b/internal/cmd/insights/queries.go new file mode 100644 index 00000000..1674dac1 --- /dev/null +++ b/internal/cmd/insights/queries.go @@ -0,0 +1,141 @@ +package insights + +import ( + "fmt" + "slices" + "strings" + + "github.com/spf13/cobra" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" +) + +// sortMetrics are the sort keys the insights API accepts that are useful for +// ranking queries from the CLI. +var sortMetrics = []string{ + "totalTime", + "count", + "errorCount", + "rowsRead", + "rowsReturned", + "rowsAffected", + "rowsReadPerReturned", + "p50Latency", + "p99Latency", + "maxLatency", + "cpuTime", + "ioTime", + "lastRun", +} + +// QueryRow is a query statistics row formatted for table output. +type QueryRow struct { + Count int64 `header:"count" json:"query_count"` + TotalTimeMs float64 `header:"total time (ms)" json:"sum_total_duration_millis"` + TimePerQryMs float64 `header:"per query (ms)" json:"time_per_query"` + P50Ms float64 `header:"p50 (ms)" json:"p50_latency"` + P99Ms float64 `header:"p99 (ms)" json:"p99_latency"` + RowsRead int64 `header:"rows read" json:"sum_rows_read"` + ReadPerReturn float64 `header:"read/returned" json:"rows_read_per_returned"` + Errors int64 `header:"errors" json:"error_count"` + LastRunAt string `header:"last run" json:"last_run_at"` + Query string `header:"query" json:"normalized_sql"` +} + +// QueriesCmd lists aggregated query statistics for a branch. +func QueriesCmd(ch *cmdutil.Helper) *cobra.Command { + var flags struct { + sort string + dir string + limit int + period string + } + + cmd := &cobra.Command{ + Use: "queries ", + Short: "List the top queries by a performance metric", + Example: ` # Queries taking the most cumulative time + pscale insights queries mydb main --org myorg + + # Most expensive read patterns (rows read vs returned) + pscale insights queries mydb main --org myorg --sort rowsReadPerReturned + + # Highest p99 latency over the last hour + pscale insights queries mydb main --org myorg --sort p99Latency --period 1h`, + Args: cmdutil.RequiredArgs("database", "branch"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database, branch := args[0], args[1] + + if !slices.Contains(sortMetrics, flags.sort) { + return fmt.Errorf("invalid --sort %q, must be one of: %s", flags.sort, strings.Join(sortMetrics, ", ")) + } + if flags.dir != "asc" && flags.dir != "desc" { + return fmt.Errorf("invalid --dir %q, must be asc or desc", flags.dir) + } + + client, err := ch.Client() + if err != nil { + return err + } + + end := ch.Printer.PrintProgress(fmt.Sprintf("Fetching query insights for %s in %s...", + printer.BoldBlue(branch), printer.BoldBlue(database))) + defer end() + + insights, err := client.QueryInsights.ListQueries(ctx, &ps.ListQueryInsightsRequest{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + }, ps.WithPerPage(flags.limit), ps.WithSort(flags.sort, flags.dir), ps.WithPeriod(flags.period)) + if err != nil { + return notFoundError(ch, err, database, branch) + } + end() + + if len(insights) == 0 && ch.Printer.Format() == printer.Human { + ch.Printer.Printf("No query statistics recorded for %s in %s.\n", + printer.BoldBlue(branch), printer.BoldBlue(database)) + return nil + } + + if ch.Printer.Format() == printer.JSON { + return ch.Printer.PrintJSON(insights) + } + return ch.Printer.PrintResource(toQueryRows(insights)) + }, + } + + cmd.Flags().StringVar(&flags.sort, "sort", "totalTime", + fmt.Sprintf("Metric to rank queries by, one of: %s", strings.Join(sortMetrics, ", "))) + cmd.Flags().StringVar(&flags.dir, "dir", "desc", "Sort direction: asc or desc") + cmd.Flags().IntVar(&flags.limit, "limit", 15, "Number of queries to return") + cmd.Flags().StringVar(&flags.period, "period", "", "Time period to aggregate over (e.g. 1h, 24h)") + + return cmd +} + +func toQueryRows(insights []*ps.QueryInsight) []*QueryRow { + rows := make([]*QueryRow, 0, len(insights)) + for _, in := range insights { + rows = append(rows, &QueryRow{ + Count: in.QueryCount, + TotalTimeMs: round2(in.SumTotalDurationMillis), + TimePerQryMs: round2(in.TimePerQuery), + P50Ms: round2(in.P50Latency), + P99Ms: round2(in.P99Latency), + RowsRead: in.SumRowsRead, + ReadPerReturn: round2(in.RowsReadPerReturned), + Errors: in.ErrorCount, + LastRunAt: in.LastRunAt.Format("2006-01-02 15:04"), + Query: truncate(in.NormalizedSQL, 80), + }) + } + return rows +} + +func round2(f float64) float64 { + return float64(int64(f*100+0.5)) / 100 +} diff --git a/internal/cmd/insights/queries_test.go b/internal/cmd/insights/queries_test.go new file mode 100644 index 00000000..e458dd9a --- /dev/null +++ b/internal/cmd/insights/queries_test.go @@ -0,0 +1,168 @@ +package insights + +import ( + "bytes" + "context" + "encoding/json" + "testing" + "time" + + "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" + + qt "github.com/frankban/quicktest" +) + +func testHelper(buf *bytes.Buffer, format printer.Format, client *ps.Client) *cmdutil.Helper { + p := printer.NewPrinter(&format) + p.SetResourceOutput(buf) + + return &cmdutil.Helper{ + Printer: p, + Config: &config.Config{Organization: "planetscale"}, + Client: func() (*ps.Client, error) { + return client, nil + }, + } +} + +func TestInsights_QueriesCmd(t *testing.T) { + c := qt.New(t) + + svc := &mock.QueryInsightsService{ + ListQueriesFn: func(ctx context.Context, req *ps.ListQueryInsightsRequest, opts ...ps.ListOption) ([]*ps.QueryInsight, error) { + c.Assert(req.Organization, qt.Equals, "planetscale") + c.Assert(req.Database, qt.Equals, "mydb") + c.Assert(req.Branch, qt.Equals, "main") + return []*ps.QueryInsight{{ + ID: "q1", + NormalizedSQL: "select * from users where id = ?", + QueryCount: 10, + SumTotalDurationMillis: 123.456, + P99Latency: 4.2, + LastRunAt: time.Date(2026, 7, 22, 23, 22, 0, 0, time.UTC), + }}, nil + }, + } + + var buf bytes.Buffer + ch := testHelper(&buf, printer.JSON, &ps.Client{QueryInsights: svc}) + + cmd := QueriesCmd(ch) + cmd.SetArgs([]string{"mydb", "main"}) + err := cmd.Execute() + + c.Assert(err, qt.IsNil) + c.Assert(svc.ListQueriesFnInvoked, qt.IsTrue) + + var out []map[string]any + c.Assert(json.Unmarshal(buf.Bytes(), &out), qt.IsNil) + c.Assert(out, qt.HasLen, 1) + c.Assert(out[0]["normalized_sql"], qt.Equals, "select * from users where id = ?") +} + +func TestInsights_QueriesCmd_InvalidSort(t *testing.T) { + c := qt.New(t) + + var buf bytes.Buffer + ch := testHelper(&buf, printer.JSON, &ps.Client{}) + + cmd := QueriesCmd(ch) + cmd.SetArgs([]string{"mydb", "main", "--sort", "bogus"}) + err := cmd.Execute() + + c.Assert(err, qt.IsNotNil) + c.Assert(err.Error(), qt.Contains, `invalid --sort "bogus"`) +} + +func TestInsights_ErrorsCmd(t *testing.T) { + c := qt.New(t) + + svc := &mock.QueryInsightsService{ + ListErrorsFn: func(ctx context.Context, req *ps.ListQueryInsightsErrorsRequest, opts ...ps.ListOption) ([]*ps.QueryInsightError, error) { + return []*ps.QueryInsightError{{ + ID: "e1", + ErrorCount: 4, + ErrorMessage: "relation \"widgets\" does not exist", + StartedAt: time.Date(2026, 7, 22, 14, 5, 0, 0, time.UTC), + }}, nil + }, + } + + var buf bytes.Buffer + ch := testHelper(&buf, printer.CSV, &ps.Client{QueryInsights: svc}) + + cmd := ErrorsCmd(ch) + cmd.SetArgs([]string{"mydb", "main"}) + err := cmd.Execute() + + c.Assert(err, qt.IsNil) + c.Assert(svc.ListErrorsFnInvoked, qt.IsTrue) + c.Assert(buf.String(), qt.Contains, "relation") +} + +func TestInsights_AnomaliesCmd(t *testing.T) { + c := qt.New(t) + + svc := &mock.QueryInsightsService{ + ListAnomaliesFn: func(ctx context.Context, req *ps.ListAnomaliesRequest, opts ...ps.ListOption) ([]*ps.Anomaly, error) { + return []*ps.Anomaly{{ + ID: "a1", + Active: true, + MinutesInViolation: 12, + PeriodStart: time.Date(2026, 7, 22, 14, 0, 0, 0, time.UTC), + }}, nil + }, + } + + var buf bytes.Buffer + ch := testHelper(&buf, printer.JSON, &ps.Client{QueryInsights: svc}) + + cmd := AnomaliesCmd(ch) + cmd.SetArgs([]string{"mydb", "main"}) + err := cmd.Execute() + + c.Assert(err, qt.IsNil) + c.Assert(svc.ListAnomaliesFnInvoked, qt.IsTrue) + + var out []map[string]any + c.Assert(json.Unmarshal(buf.Bytes(), &out), qt.IsNil) + c.Assert(out[0]["id"], qt.Equals, "a1") + c.Assert(out[0]["active"], qt.Equals, true) +} + +func TestInsights_RecommendationsCmd(t *testing.T) { + c := qt.New(t) + + svc := &mock.SchemaRecommendationService{ + ListFn: func(ctx context.Context, req *ps.ListSchemaRecommendationsRequest, opts ...ps.ListOption) ([]*ps.SchemaRecommendation, error) { + c.Assert(req.Database, qt.Equals, "mydb") + return []*ps.SchemaRecommendation{{ + Number: 1, + State: "open", + RecommendationType: "duplicate_index", + Table: "users", + Title: "Drop duplicate index idx_email", + DDLStatement: "ALTER TABLE users DROP INDEX idx_email", + }}, nil + }, + } + + var buf bytes.Buffer + ch := testHelper(&buf, printer.JSON, &ps.Client{SchemaRecommendations: svc}) + + cmd := RecommendationsCmd(ch) + cmd.SetArgs([]string{"mydb"}) + err := cmd.Execute() + + c.Assert(err, qt.IsNil) + c.Assert(svc.ListFnInvoked, qt.IsTrue) + + var out []map[string]any + c.Assert(json.Unmarshal(buf.Bytes(), &out), qt.IsNil) + c.Assert(out[0]["recommendation_type"], qt.Equals, "duplicate_index") + c.Assert(out[0]["ddl_statement"], qt.Equals, "ALTER TABLE users DROP INDEX idx_email") +} diff --git a/internal/cmd/insights/recommendations.go b/internal/cmd/insights/recommendations.go new file mode 100644 index 00000000..45ed5991 --- /dev/null +++ b/internal/cmd/insights/recommendations.go @@ -0,0 +1,82 @@ +package insights + +import ( + "fmt" + + "github.com/spf13/cobra" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" +) + +// RecommendationRow is a schema recommendation formatted for table output. +type RecommendationRow struct { + Number int `header:"number" json:"number"` + State string `header:"state" json:"state"` + Type string `header:"type" json:"recommendation_type"` + Table string `header:"table" json:"table_name"` + Keyspace string `header:"keyspace" json:"keyspace,omitempty"` + Title string `header:"title" json:"title"` +} + +// RecommendationsCmd lists schema recommendations for a database. +func RecommendationsCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "recommendations ", + Short: "List schema recommendations (unused/duplicate indexes, bloat, missing indexes)", + Long: `List PlanetScale's schema recommendations for a database: unused tables and +indexes, duplicate indexes, bloated tables and indexes, missing indexes +derived from production query patterns, and sequence overflow risks. Each +recommendation includes ready-to-apply DDL (shown with --format json).`, + Aliases: []string{"recommendation"}, + Args: cmdutil.RequiredArgs("database"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database := args[0] + + client, err := ch.Client() + if err != nil { + return err + } + + end := ch.Printer.PrintProgress(fmt.Sprintf("Fetching schema recommendations for %s...", + printer.BoldBlue(database))) + defer end() + + recommendations, err := client.SchemaRecommendations.List(ctx, &ps.ListSchemaRecommendationsRequest{ + Organization: ch.Config.Organization, + Database: database, + }) + if err != nil { + return notFoundError(ch, err, database, "") + } + end() + + if len(recommendations) == 0 && ch.Printer.Format() == printer.Human { + ch.Printer.Printf("No schema recommendations for %s. Your schema looks good!\n", + printer.BoldBlue(database)) + return nil + } + + if ch.Printer.Format() == printer.JSON { + return ch.Printer.PrintJSON(recommendations) + } + + rows := make([]*RecommendationRow, 0, len(recommendations)) + for _, r := range recommendations { + rows = append(rows, &RecommendationRow{ + Number: r.Number, + State: r.State, + Type: r.RecommendationType, + Table: r.Table, + Keyspace: r.Keyspace, + Title: truncate(r.Title, 80), + }) + } + return ch.Printer.PrintResource(rows) + }, + } + + return cmd +} diff --git a/internal/cmd/inspect/checks.go b/internal/cmd/inspect/checks.go new file mode 100644 index 00000000..64f3c83a --- /dev/null +++ b/internal/cmd/inspect/checks.go @@ -0,0 +1,624 @@ +package inspect + +// engineSQL is one engine's implementation of a check: a single read-only +// diagnostic query, plus an optional PostgreSQL extension requirement. +type engineSQL struct { + // SQL is the read-only query to run. Result sets must be bounded (LIMIT) + // so output stays safe for agents and terminals. + SQL string + // RequiresExtension names a PostgreSQL extension the query depends on. + // When set, the runner checks pg_extension first and fails with an + // actionable hint if it isn't installed. + RequiresExtension string +} + +// check is a named diagnostic. A nil engine entry means the check doesn't +// apply to that engine; Hint explains what to use instead. +type check struct { + Name string + Short string + EmptyMessage string + MySQL *engineSQL + Postgres *engineSQL + // MySQLHint / PostgresHint are shown when the check isn't available for + // that engine. + MySQLHint string + PostgresHint string + // NextSteps are related commands (usually pscale insights) that provide + // server-side, traffic-aware analysis of the same problem. They are shown + // after results and in JSON output so agents know to cross-reference. + // Occurrences of and are replaced with the real args. + NextSteps []string +} + +// mysqlSystemSchemas excludes MySQL system schemas and Vitess internal +// schemas from results. +const mysqlSystemSchemas = `('information_schema', 'performance_schema', 'mysql', 'sys', '_vt')` + +var checks = []check{ + { + Name: "table-sizes", + Short: "Tables by total size, largest first", + EmptyMessage: "No user tables found.", + MySQL: &engineSQL{ + SQL: ` + SELECT + table_schema AS ` + "`schema`" + `, + table_name AS name, + ROUND((data_length + index_length) / 1024 / 1024, 2) AS size_mb, + table_rows AS approx_rows + FROM information_schema.tables + WHERE table_schema NOT IN ` + mysqlSystemSchemas + ` + AND table_type = 'BASE TABLE' + ORDER BY data_length + index_length DESC + LIMIT 25;`, + }, + Postgres: &engineSQL{ + // pg_table_size (heap + TOAST + FSM + VM, no indexes) — index + // storage is itemized by the index-sizes check, so including it + // here would double-count. Partitions roll up into their root so + // a partitioned table shows as one row; materialized views are + // included since they also consume storage. + SQL: ` + SELECT + schema, + name, + type, + pg_size_pretty(table_bytes) AS size + FROM ( + SELECT + n.nspname AS schema, + root.relname AS name, + CASE root.relkind + WHEN 'm' THEN 'matview' + WHEN 'p' THEN 'partitioned table' + ELSE 'table' + END AS type, + SUM(pg_table_size(c.oid)) AS table_bytes + FROM pg_class c + JOIN pg_class root ON root.oid = COALESCE(pg_partition_root(c.oid), c.oid) + JOIN pg_namespace n ON n.oid = root.relnamespace + WHERE c.relkind IN ('r', 'm', 'p') + AND n.nspname NOT IN ('pg_catalog', 'information_schema') + AND n.nspname !~ '^pg_toast' + GROUP BY n.nspname, root.relname, root.relkind + ) sizes + ORDER BY table_bytes DESC + LIMIT 25;`, + }, + }, + { + Name: "index-sizes", + Short: "Indexes by size, largest first", + EmptyMessage: "No indexes found.", + MySQL: &engineSQL{ + SQL: ` + SELECT + database_name AS ` + "`schema`" + `, + table_name AS ` + "`table`" + `, + index_name, + ROUND(stat_value * @@innodb_page_size / 1024 / 1024, 2) AS size_mb + FROM mysql.innodb_index_stats + WHERE stat_name = 'size' + AND database_name NOT IN ('mysql', 'sys', '_vt') + ORDER BY stat_value DESC + LIMIT 25;`, + }, + Postgres: &engineSQL{ + SQL: ` + SELECT + n.nspname AS schema, + t.relname || '.' || c.relname AS name, + pg_size_pretty(pg_relation_size(c.oid)) AS size + FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + JOIN pg_class t ON t.oid = i.indrelid + LEFT JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') + AND n.nspname !~ '^pg_toast' + AND c.relkind = 'i' + ORDER BY pg_relation_size(c.oid) DESC + LIMIT 25;`, + }, + }, + { + Name: "unused-indexes", + NextSteps: []string{"pscale insights recommendations "}, + Short: "Indexes with little or no use — removal candidates", + EmptyMessage: "No unused indexes detected.", + MySQL: &engineSQL{ + SQL: ` + SELECT + object_schema AS ` + "`schema`" + `, + object_name AS ` + "`table`" + `, + index_name + FROM sys.schema_unused_indexes + WHERE object_schema NOT IN ` + mysqlSystemSchemas + ` + ORDER BY object_schema, object_name + LIMIT 50;`, + }, + Postgres: &engineSQL{ + SQL: ` + SELECT + s.schemaname || '.' || s.relname AS "table", + s.indexrelname AS index, + pg_size_pretty(pg_relation_size(s.indexrelid)) AS index_size, + s.idx_scan AS index_scans + FROM pg_stat_user_indexes s + JOIN pg_index i ON i.indexrelid = s.indexrelid + WHERE NOT i.indisunique + AND NOT i.indisprimary + AND s.idx_scan < 50 + ORDER BY pg_relation_size(s.indexrelid) DESC, s.idx_scan ASC + LIMIT 50;`, + }, + }, + { + Name: "redundant-indexes", + NextSteps: []string{"pscale insights recommendations "}, + Short: "Indexes made redundant by another index", + EmptyMessage: "No redundant indexes detected.", + MySQL: &engineSQL{ + SQL: ` + SELECT + table_schema AS ` + "`schema`" + `, + table_name AS ` + "`table`" + `, + redundant_index_name, + redundant_index_columns, + dominant_index_name, + dominant_index_columns + FROM sys.schema_redundant_indexes + WHERE table_schema NOT IN ` + mysqlSystemSchemas + ` + LIMIT 50;`, + }, + PostgresHint: "Redundant-index detection for PostgreSQL is served by schema recommendations: pscale insights recommendations ", + }, + { + Name: "invalid-indexes", + Short: "Invalid indexes left over from failed concurrent index builds", + EmptyMessage: "No invalid indexes found.", + MySQLHint: "Invalid indexes are a PostgreSQL concept (left over from a failed CREATE INDEX CONCURRENTLY).", + Postgres: &engineSQL{ + SQL: ` + SELECT + n.nspname AS schema, + t.relname || '.' || c.relname AS name, + pg_size_pretty(pg_relation_size(c.oid)) AS size, + i.indisready AS ready + FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + JOIN pg_class t ON t.oid = i.indrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE NOT i.indisvalid + AND n.nspname NOT IN ('pg_catalog', 'information_schema') + AND n.nspname !~ '^pg_toast' + ORDER BY pg_relation_size(c.oid) DESC + LIMIT 50;`, + }, + }, + { + Name: "seq-scans", + NextSteps: []string{"pscale insights queries --sort rowsReadPerReturned"}, + Short: "Tables receiving full-table scans", + EmptyMessage: "No full-table scans recorded.", + MySQL: &engineSQL{ + SQL: ` + SELECT + object_schema AS ` + "`schema`" + `, + object_name AS name, + rows_full_scanned, + latency + FROM sys.schema_tables_with_full_table_scans + WHERE object_schema NOT IN ` + mysqlSystemSchemas + ` + ORDER BY rows_full_scanned DESC + LIMIT 25;`, + }, + Postgres: &engineSQL{ + SQL: ` + SELECT + schemaname AS schema, + relname AS name, + seq_scan AS count + FROM pg_stat_user_tables + WHERE seq_scan > 0 + ORDER BY seq_scan DESC + LIMIT 25;`, + }, + }, + { + Name: "long-running-queries", + NextSteps: []string{"pscale insights queries --sort p99Latency"}, + Short: "Queries running longer than 5 minutes", + EmptyMessage: "No long-running queries.", + MySQL: &engineSQL{ + SQL: ` + SELECT + id, + user, + db, + time AS seconds, + state, + LEFT(COALESCE(info, ''), 120) AS query + FROM information_schema.processlist + WHERE command NOT IN ('Sleep', 'Binlog Dump', 'Binlog Dump GTID', 'Daemon') + AND user NOT IN ('vt_repl', 'event_scheduler') + AND info IS NOT NULL + AND time > 300 + ORDER BY time DESC + LIMIT 100;`, + }, + Postgres: &engineSQL{ + // walsenders are excluded: replication connections stay "active" + // forever and would always appear here. + SQL: ` + SELECT + pid, + (now() - query_start)::text AS duration, + state, + left(query, 200) AS query + FROM pg_stat_activity + WHERE state <> 'idle' + AND backend_type <> 'walsender' + AND query NOT ILIKE '%pg_stat_activity%' + AND now() - query_start > interval '5 minutes' + ORDER BY now() - query_start DESC + LIMIT 100;`, + }, + }, + { + Name: "locks", + Short: "Blocking locks and the sessions stuck behind them", + EmptyMessage: "No blocking locks.", + MySQL: &engineSQL{ + SQL: ` + SELECT + waiting_pid, + LEFT(COALESCE(waiting_query, ''), 80) AS waiting_query, + blocking_pid, + LEFT(COALESCE(blocking_query, ''), 80) AS blocking_query, + wait_age, + locked_table + FROM sys.innodb_lock_waits + LIMIT 100;`, + }, + Postgres: &engineSQL{ + // Reports only the roots of blocking trees (via pg_blocking_pids) + // with a count of sessions stuck behind each, instead of every + // granted lock — most locks are routine and non-blocking. + SQL: ` + WITH RECURSIVE waiters AS ( + SELECT a.pid AS blocked_pid, b.pid AS blocking_pid + FROM pg_stat_activity a + CROSS JOIN LATERAL unnest(pg_blocking_pids(a.pid)) AS b(pid) + WHERE a.wait_event_type = 'Lock' + ), + locks AS MATERIALIZED ( + SELECT * FROM pg_locks + ), + chain AS ( + SELECT blocked_pid, blocking_pid AS root_pid + FROM waiters + WHERE blocking_pid NOT IN (SELECT blocked_pid FROM waiters) + UNION + SELECT w.blocked_pid, c.root_pid + FROM waiters w + JOIN chain c ON c.blocked_pid = w.blocking_pid + ) + SELECT + a.pid, + a.state, + string_agg(DISTINCT wc.relname, ', ') AS relation, + string_agg(DISTINCT hl.mode, ', ') AS lock_mode, + string_agg(DISTINCT hl.locktype, ', ') AS lock_type, + (now() - a.query_start)::text AS age, + count(DISTINCT ch.blocked_pid) AS blocked_count, + left(a.query, 200) AS query + FROM chain ch + JOIN pg_stat_activity a ON a.pid = ch.root_pid + LEFT JOIN locks wl ON wl.pid = ch.blocked_pid AND NOT wl.granted + LEFT JOIN pg_class wc ON wc.oid = wl.relation + LEFT JOIN locks hl ON hl.pid = a.pid AND hl.granted + AND hl.locktype = wl.locktype + AND hl.database IS NOT DISTINCT FROM wl.database + AND hl.relation IS NOT DISTINCT FROM wl.relation + AND hl.page IS NOT DISTINCT FROM wl.page + AND hl.tuple IS NOT DISTINCT FROM wl.tuple + AND hl.virtualxid IS NOT DISTINCT FROM wl.virtualxid + AND hl.transactionid IS NOT DISTINCT FROM wl.transactionid + AND hl.classid IS NOT DISTINCT FROM wl.classid + AND hl.objid IS NOT DISTINCT FROM wl.objid + AND hl.objsubid IS NOT DISTINCT FROM wl.objsubid + GROUP BY a.pid, a.state, a.query_start, a.query + ORDER BY blocked_count DESC, a.query_start + LIMIT 100;`, + }, + }, + { + Name: "outliers", + NextSteps: []string{"pscale insights queries --sort totalTime"}, + Short: "Queries by cumulative execution time (needs pg_stat_statements)", + EmptyMessage: "No statements recorded yet.", + MySQLHint: "Query timing for MySQL is served by insights: pscale insights queries --sort totalTime", + Postgres: &engineSQL{ + RequiresExtension: "pg_stat_statements", + SQL: ` + SELECT + (interval '1 millisecond' * total_exec_time)::text AS total_exec_time, + to_char( + (total_exec_time / nullif(sum(total_exec_time) OVER (), 0)) * 100, + 'FM990D0' + ) || '%' AS prop_exec_time, + calls AS ncalls, + query + FROM pg_stat_statements s + JOIN pg_database d ON d.oid = s.dbid + WHERE d.datname = current_database() + ORDER BY total_exec_time DESC + LIMIT 25;`, + }, + }, + { + Name: "calls", + NextSteps: []string{"pscale insights queries --sort count"}, + Short: "Most frequently called queries (needs pg_stat_statements)", + EmptyMessage: "No statements recorded yet.", + MySQLHint: "Query frequency for MySQL is served by insights: pscale insights queries --sort count", + Postgres: &engineSQL{ + RequiresExtension: "pg_stat_statements", + SQL: ` + SELECT + calls AS ncalls, + (interval '1 millisecond' * total_exec_time)::text AS total_exec_time, + to_char( + (calls / nullif(sum(calls) OVER (), 0)::float) * 100, + 'FM990D0' + ) || '%' AS prop_calls, + query + FROM pg_stat_statements s + JOIN pg_database d ON d.oid = s.dbid + WHERE d.datname = current_database() + ORDER BY calls DESC + LIMIT 25;`, + }, + }, + { + Name: "bloat", + NextSteps: []string{"pscale insights recommendations "}, + Short: "Wasted space: estimated bloat (PostgreSQL) or fragmentation (MySQL)", + EmptyMessage: "No bloat detected.", + MySQL: &engineSQL{ + SQL: ` + SELECT + table_schema AS ` + "`schema`" + `, + table_name AS name, + ROUND(data_free / 1024 / 1024, 2) AS free_mb, + ROUND((data_length + index_length) / 1024 / 1024, 2) AS size_mb, + ROUND(data_free / NULLIF(data_length + index_length + data_free, 0) * 100, 1) AS free_pct + FROM information_schema.tables + WHERE table_schema NOT IN ` + mysqlSystemSchemas + ` + AND data_free > 0 + ORDER BY data_free DESC + LIMIT 25;`, + }, + Postgres: &engineSQL{ + // Statistical estimate covering both tables and indexes. The + // index estimate is rough: it assumes the index holds all table + // columns. + SQL: ` + WITH constants AS ( + SELECT current_setting('block_size')::numeric AS bs, 23 AS hdr, 8 AS ma + ), + bloat_info AS ( + SELECT + ma, bs, schemaname, tablename, + (datawidth + (hdr + ma - (CASE WHEN hdr % ma = 0 THEN ma ELSE hdr % ma END)))::numeric AS datahdr, + (maxfracsum * (nullhdr + ma - (CASE WHEN nullhdr % ma = 0 THEN ma ELSE nullhdr % ma END))) AS nullhdr2 + FROM ( + SELECT + schemaname, tablename, hdr, ma, bs, + SUM((1 - null_frac) * avg_width) AS datawidth, + MAX(null_frac) AS maxfracsum, + hdr + 1 + count(*) FILTER (WHERE null_frac <> 0) / 8 AS nullhdr + FROM pg_stats s, constants + GROUP BY schemaname, tablename, hdr, ma, bs + ) AS foo + ), + table_bloat AS ( + SELECT + schemaname, tablename, cc.relpages, bs, + CEIL((cc.reltuples * ((datahdr + ma + - (CASE WHEN datahdr % ma = 0 THEN ma ELSE datahdr % ma END)) + + nullhdr2 + 4)) / (bs - 20::float)) AS otta + FROM bloat_info + JOIN pg_class cc ON cc.relname = bloat_info.tablename + JOIN pg_namespace nn ON cc.relnamespace = nn.oid + AND nn.nspname = bloat_info.schemaname + AND nn.nspname NOT IN ('information_schema', 'pg_catalog') + WHERE cc.relkind = 'r' + ), + index_bloat AS ( + SELECT + bloat_info.schemaname, bloat_info.tablename, bs, + c2.relname AS iname, + c2.relpages AS ipages, + COALESCE(CEIL((c2.reltuples * (datahdr - 12)) / (bs - 20::float)), 0) AS iotta + FROM bloat_info + JOIN pg_class cc ON cc.relname = bloat_info.tablename + JOIN pg_namespace nn ON cc.relnamespace = nn.oid + AND nn.nspname = bloat_info.schemaname + AND nn.nspname NOT IN ('information_schema', 'pg_catalog') + JOIN pg_index i ON i.indrelid = cc.oid + JOIN pg_class c2 ON c2.oid = i.indexrelid + WHERE cc.relkind = 'r' + ) + SELECT + type, schema, object_name, + pg_size_pretty(size) AS size, + bloat_pct, + pg_size_pretty(raw_waste) AS waste + FROM ( + SELECT + 'table' AS type, + schemaname AS schema, + tablename AS object_name, + (bs * relpages)::bigint AS size, + round(CASE WHEN relpages = 0 OR relpages < otta THEN 0.0 + ELSE 100 * (relpages - otta)::numeric / relpages + END, 1) AS bloat_pct, + (CASE WHEN relpages < otta THEN 0 ELSE bs * (relpages - otta) END)::bigint AS raw_waste + FROM table_bloat + UNION ALL + SELECT + 'index' AS type, + schemaname AS schema, + tablename || '::' || iname AS object_name, + (bs * ipages)::bigint AS size, + round(CASE WHEN ipages = 0 OR ipages < iotta THEN 0.0 + ELSE 100 * (ipages - iotta)::numeric / ipages + END, 1) AS bloat_pct, + (CASE WHEN ipages < iotta THEN 0 ELSE bs * (ipages - iotta) END)::bigint AS raw_waste + FROM index_bloat + ) summary + ORDER BY raw_waste DESC, bloat_pct DESC + LIMIT 25;`, + }, + }, + { + Name: "vacuum-stats", + Short: "Autovacuum and autoanalyze health: last runs, dead tuples, thresholds", + EmptyMessage: "No user tables found.", + MySQLHint: "Vacuum is a PostgreSQL concept; for MySQL fragmentation see: pscale inspect bloat", + Postgres: &engineSQL{ + // Thresholds honor per-table reloptions (including + // autovacuum_enabled=false, reported as "disabled"), not just the + // global settings. + SQL: ` + SELECT + psut.schemaname AS schema, + psut.relname AS "table", + to_char(psut.last_vacuum, 'YYYY-MM-DD HH24:MI') AS last_vacuum, + to_char(psut.last_autovacuum, 'YYYY-MM-DD HH24:MI') AS last_autovacuum, + to_char(psut.last_analyze, 'YYYY-MM-DD HH24:MI') AS last_analyze, + to_char(psut.last_autoanalyze, 'YYYY-MM-DD HH24:MI') AS last_autoanalyze, + c.reltuples::bigint AS rowcount, + psut.n_dead_tup AS dead_rowcount, + psut.n_mod_since_analyze AS mods_since_analyze, + CASE + WHEN opts.enabled = 'false' THEN 'disabled' + WHEN opts.vac_threshold + opts.vac_scale * GREATEST(c.reltuples, 0) + < psut.n_dead_tup THEN 'yes' + ELSE 'no' + END AS expect_autovacuum, + CASE + WHEN opts.enabled = 'false' THEN 'disabled' + WHEN opts.an_threshold + opts.an_scale * GREATEST(c.reltuples, 0) + < psut.n_mod_since_analyze THEN 'yes' + ELSE 'no' + END AS expect_autoanalyze + FROM pg_stat_user_tables psut + JOIN pg_class c ON c.oid = psut.relid + CROSS JOIN LATERAL ( + SELECT + COALESCE((SELECT split_part(o, '=', 2)::bigint + FROM unnest(c.reloptions) AS o + WHERE o LIKE 'autovacuum_vacuum_threshold=%'), + current_setting('autovacuum_vacuum_threshold')::bigint) AS vac_threshold, + COALESCE((SELECT split_part(o, '=', 2)::numeric + FROM unnest(c.reloptions) AS o + WHERE o LIKE 'autovacuum_vacuum_scale_factor=%'), + current_setting('autovacuum_vacuum_scale_factor')::numeric) AS vac_scale, + COALESCE((SELECT split_part(o, '=', 2)::bigint + FROM unnest(c.reloptions) AS o + WHERE o LIKE 'autovacuum_analyze_threshold=%'), + current_setting('autovacuum_analyze_threshold')::bigint) AS an_threshold, + COALESCE((SELECT split_part(o, '=', 2)::numeric + FROM unnest(c.reloptions) AS o + WHERE o LIKE 'autovacuum_analyze_scale_factor=%'), + current_setting('autovacuum_analyze_scale_factor')::numeric) AS an_scale, + COALESCE((SELECT split_part(o, '=', 2) + FROM unnest(c.reloptions) AS o + WHERE o LIKE 'autovacuum_enabled=%'), + 'true') AS enabled + ) opts + ORDER BY psut.n_dead_tup DESC + LIMIT 25;`, + }, + }, + { + Name: "replication-slots", + Short: "Replication slots: status, WAL retention, and lag", + EmptyMessage: "No replication slots found.", + MySQLHint: "Replication slots are a PostgreSQL concept; for Vitess workflows see: pscale workflow list", + Postgres: &engineSQL{ + // retained_wal_size (since restart_lsn) and unconfirmed_wal_size + // (since confirmed_flush_lsn) measure different failure modes; + // safe_wal_size shows headroom before the slot is invalidated by + // max_slot_wal_keep_size. + SQL: ` + SELECT + s.slot_name, + s.slot_type, + s.database, + CASE + WHEN r.state IS NOT NULL THEN r.state + WHEN s.active THEN 'active (no walsender)' + ELSE 'inactive' + END AS status, + s.temporary, + s.wal_status, + pg_size_pretty(s.safe_wal_size) AS safe_wal_size, + pg_size_pretty( + pg_wal_lsn_diff( + CASE WHEN pg_is_in_recovery() THEN pg_last_wal_receive_lsn() + ELSE pg_current_wal_lsn() END, + s.restart_lsn + ) + ) AS retained_wal_size, + CASE + WHEN s.slot_type = 'logical' THEN + pg_size_pretty( + pg_wal_lsn_diff( + CASE WHEN pg_is_in_recovery() THEN pg_last_wal_receive_lsn() + ELSE pg_current_wal_lsn() END, + s.confirmed_flush_lsn + ) + ) + ELSE NULL + END AS unconfirmed_wal_size, + r.replay_lag AS lag_time + FROM pg_replication_slots s + LEFT JOIN pg_stat_replication r ON r.pid = s.active_pid + ORDER BY pg_wal_lsn_diff( + CASE WHEN pg_is_in_recovery() THEN pg_last_wal_receive_lsn() + ELSE pg_current_wal_lsn() END, + s.restart_lsn + ) DESC NULLS LAST + LIMIT 100;`, + }, + }, + { + Name: "subscriptions", + Short: "Per-table logical replication progress on this subscriber", + EmptyMessage: "No subscriptions found on this database.", + MySQLHint: "Subscriptions are a PostgreSQL concept; for imports see: pscale data-imports get", + Postgres: &engineSQL{ + SQL: ` + SELECT + sub.subname AS subscription, + sr.srrelid::regclass::text AS table_name, + CASE sr.srsubstate + WHEN 'i' THEN 'initialize' + WHEN 'd' THEN 'data being copied' + WHEN 'f' THEN 'finished table copy' + WHEN 's' THEN 'synchronized' + WHEN 'r' THEN 'ready' + ELSE sr.srsubstate::text + END AS status, + sr.srsublsn AS lsn + FROM pg_subscription_rel sr + JOIN pg_subscription sub ON sub.oid = sr.srsubid + ORDER BY sub.subname, table_name + LIMIT 100;`, + }, + }, +} diff --git a/internal/cmd/inspect/checks_test.go b/internal/cmd/inspect/checks_test.go new file mode 100644 index 00000000..22945b50 --- /dev/null +++ b/internal/cmd/inspect/checks_test.go @@ -0,0 +1,86 @@ +package inspect + +import ( + "regexp" + "strings" + "testing" + + qt "github.com/frankban/quicktest" +) + +var checkNameRE = regexp.MustCompile(`^[a-z][a-z0-9-]*$`) + +// TestCheckCatalog validates the invariants every check must hold: read-only, +// bounded SQL, an implementation or an alternative hint per engine, and +// well-formed metadata. +func TestCheckCatalog(t *testing.T) { + c := qt.New(t) + + seen := map[string]bool{} + for _, chk := range checks { + c.Assert(checkNameRE.MatchString(chk.Name), qt.IsTrue, qt.Commentf("check name %q must be kebab-case", chk.Name)) + c.Assert(seen[chk.Name], qt.IsFalse, qt.Commentf("duplicate check name %q", chk.Name)) + seen[chk.Name] = true + + c.Assert(chk.Short, qt.Not(qt.Equals), "", qt.Commentf("%s: missing Short", chk.Name)) + c.Assert(chk.EmptyMessage, qt.Not(qt.Equals), "", qt.Commentf("%s: missing EmptyMessage", chk.Name)) + c.Assert(chk.MySQL != nil || chk.Postgres != nil, qt.IsTrue, + qt.Commentf("%s: no engine implementation", chk.Name)) + + if chk.MySQL == nil { + c.Assert(chk.MySQLHint, qt.Not(qt.Equals), "", qt.Commentf("%s: MySQL unsupported but no hint", chk.Name)) + } + if chk.Postgres == nil { + c.Assert(chk.PostgresHint, qt.Not(qt.Equals), "", qt.Commentf("%s: Postgres unsupported but no hint", chk.Name)) + } + + for engine, impl := range map[string]*engineSQL{"mysql": chk.MySQL, "postgres": chk.Postgres} { + if impl == nil { + continue + } + sql := strings.ToUpper(strings.TrimSpace(impl.SQL)) + c.Assert(strings.HasPrefix(sql, "SELECT") || strings.HasPrefix(sql, "WITH"), qt.IsTrue, + qt.Commentf("%s/%s: diagnostic SQL must be a read-only SELECT", chk.Name, engine)) + c.Assert(strings.Contains(sql, "LIMIT"), qt.IsTrue, + qt.Commentf("%s/%s: diagnostic SQL must bound its result set with LIMIT", chk.Name, engine)) + c.Assert(strings.Count(impl.SQL, ";"), qt.Equals, 1, + qt.Commentf("%s/%s: diagnostic SQL must be a single statement", chk.Name, engine)) + // table-sizes must exclude index storage (pg_table_size, not + // pg_total_relation_size): indexes are itemized by index-sizes, + // so counting them here would double-count across checks. + if chk.Name == "table-sizes" && engine == "postgres" { + c.Assert(strings.Contains(impl.SQL, "pg_table_size"), qt.IsTrue) + c.Assert(strings.Contains(impl.SQL, "pg_total_relation_size"), qt.IsFalse) + } + } + + for _, step := range chk.NextSteps { + c.Assert(strings.HasPrefix(step, "pscale "), qt.IsTrue, + qt.Commentf("%s: next step %q must be a pscale command", chk.Name, step)) + } + } +} + +func TestFormatNextSteps(t *testing.T) { + c := qt.New(t) + + steps := formatNextSteps([]string{ + "pscale insights queries --sort totalTime", + "pscale insights recommendations ", + }, "myorg", "mydb", "main") + + c.Assert(steps, qt.DeepEquals, []string{ + "pscale insights queries mydb main --sort totalTime --org myorg --format json", + "pscale insights recommendations mydb --org myorg --format json", + }) +} + +func TestFormatValue(t *testing.T) { + c := qt.New(t) + + c.Assert(formatValue(nil), qt.Equals, "") + c.Assert(formatValue("plain"), qt.Equals, "plain") + c.Assert(formatValue("multi\n line\tsql"), qt.Equals, "multi line sql") + c.Assert(formatValue(int64(42)), qt.Equals, "42") + c.Assert(formatValue(1.5), qt.Equals, "1.5") +} diff --git a/internal/cmd/inspect/inspect.go b/internal/cmd/inspect/inspect.go new file mode 100644 index 00000000..e9cea40f --- /dev/null +++ b/internal/cmd/inspect/inspect.go @@ -0,0 +1,413 @@ +package inspect + +import ( + "context" + "encoding/csv" + "fmt" + "io" + "strings" + "text/tabwriter" + "time" + + "github.com/spf13/cobra" + + "github.com/planetscale/cli/internal/cmdutil" + "github.com/planetscale/cli/internal/printer" + "github.com/planetscale/cli/internal/sqlquery" +) + +// checkTimeout bounds a single diagnostic query so one slow check can't hang +// the command. +const checkTimeout = 30 * time.Second + +type inspectFlags struct { + keyspace string + postgresDB string + role string + replica bool +} + +// InspectCmd runs read-only diagnostic checks against a database branch. +func InspectCmd(ch *cmdutil.Helper) *cobra.Command { + flags := &inspectFlags{} + + cmd := &cobra.Command{ + Use: "inspect ", + Short: "Run read-only diagnostic checks against a database branch", + Long: `Run read-only diagnostic checks against a database branch, using the same +ephemeral credentials as pscale sql. Each check is a single bounded query +against the engine's statistics tables. Nothing is written to the database. + +Checks adapt to the database engine: MySQL (Vitess) checks read +information_schema, mysql, and sys; PostgreSQL checks read pg_catalog and +pg_stat views. Checks that don't apply to an engine explain what to use +instead. + +On sharded Vitess databases, statistics reflect one shard's MySQL instance +per run. Pass --keyspace to pick the keyspace, or target an exact shard with +--keyspace 'mykeyspace/-80' (enumerate shards with SHOW VITESS_SHARDS via +pscale sql). Databases can have hundreds of shards, so no check fans out +across shards automatically. + +On PostgreSQL, statistics are scoped to one database. Pass --dbname to target +the database your application uses (defaults to postgres). + +For server-side, traffic-aware analysis (slow queries, schema recommendations, +anomalies), see pscale insights.`, + PersistentPreRunE: cmdutil.CheckAuthentication(ch.Config), + } + + cmd.PersistentFlags().StringVar(&ch.Config.Organization, "org", ch.Config.Organization, + "The organization for the current user") + cmd.MarkPersistentFlagRequired("org") // nolint:errcheck + cmd.PersistentFlags().StringVar(&flags.keyspace, "keyspace", "", + "Vitess keyspace to inspect, optionally with a shard and tablet type (e.g. mykeyspace, mykeyspace/-80, mykeyspace/-80@replica). List shards with: pscale sql --query \"SHOW VITESS_SHARDS\". Defaults to @primary.") + cmd.PersistentFlags().StringVar(&flags.postgresDB, "dbname", "postgres", + "PostgreSQL database name to inspect") + cmd.PersistentFlags().StringVar(&flags.role, "role", "", + "Access role for the ephemeral credentials: reader, writer, readwriter, or admin. Defaults to reader. On PostgreSQL, the reader role may lack CONNECT on non-default databases; use --role admin if connecting with --dbname fails.") + cmd.PersistentFlags().BoolVar(&flags.replica, "replica", false, + "Run checks against a replica instead of the primary") + + for _, c := range checks { + cmd.AddCommand(checkCmd(ch, c, flags)) + } + cmd.AddCommand(allCmd(ch, flags)) + + return cmd +} + +// CheckResult is one check's outcome, printable in all output formats. +type CheckResult struct { + Check string `json:"check"` + Database string `json:"database"` + Branch string `json:"branch"` + Columns []string `json:"columns,omitempty"` + Rows []map[string]any `json:"rows"` + RowCount int `json:"row_count"` + Skipped string `json:"skipped,omitempty"` + // NextSteps points at the pscale insights commands that analyze the same + // problem from server-side production traffic data. + NextSteps []string `json:"next_steps,omitempty"` +} + +// Report is the combined output of inspect all. +type Report struct { + Database string `json:"database"` + Branch string `json:"branch"` + Results []*CheckResult `json:"results"` + // NextSteps recommends the server-side analysis commands that complement + // these connection-level checks. + NextSteps []string `json:"next_steps"` +} + +func insightsNextSteps(organization, database, branch string) []string { + return formatNextSteps([]string{ + "pscale insights queries ", + "pscale insights errors ", + "pscale insights anomalies ", + "pscale insights recommendations ", + }, organization, database, branch) +} + +func formatNextSteps(steps []string, organization, database, branch string) []string { + out := make([]string, 0, len(steps)) + for _, step := range steps { + step = strings.ReplaceAll(step, "", database) + step = strings.ReplaceAll(step, "", branch) + out = append(out, fmt.Sprintf("%s --org %s --format json", step, organization)) + } + return out +} + +func substituteResourceNames(s, database, branch string) string { + s = strings.ReplaceAll(s, "", database) + return strings.ReplaceAll(s, "", branch) +} + +func checkCmd(ch *cmdutil.Helper, c check, flags *inspectFlags) *cobra.Command { + return &cobra.Command{ + Use: c.Name + " ", + Short: c.Short, + Args: cmdutil.RequiredArgs("database", "branch"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database, branch := args[0], args[1] + + end := ch.Printer.PrintProgress(fmt.Sprintf("Inspecting %s in %s...", + printer.BoldBlue(branch), printer.BoldBlue(database))) + defer end() + + sess, err := newSession(ctx, ch, database, branch, flags) + if err != nil { + return cmdutil.HandleError(err) + } + defer sess.Close() + + result, err := runCheck(ctx, sess, c, ch.Config.Organization, database, branch) + if err != nil { + return cmdutil.HandleError(err) + } + end() + + if result.Skipped != "" && ch.Printer.Format() == printer.Human { + ch.Printer.Printf("%s\n", result.Skipped) + printNextSteps(ch, result.NextSteps) + return nil + } + + if err := printCheck(ch, c, result); err != nil { + return err + } + if ch.Printer.Format() == printer.Human { + printNextSteps(ch, result.NextSteps) + } + return nil + }, + } +} + +func allCmd(ch *cmdutil.Helper, flags *inspectFlags) *cobra.Command { + return &cobra.Command{ + Use: "all ", + Short: "Run every applicable check and print a combined report", + Args: cmdutil.RequiredArgs("database", "branch"), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + database, branch := args[0], args[1] + + end := ch.Printer.PrintProgress(fmt.Sprintf("Inspecting %s in %s...", + printer.BoldBlue(branch), printer.BoldBlue(database))) + defer end() + + sess, err := newSession(ctx, ch, database, branch, flags) + if err != nil { + return cmdutil.HandleError(err) + } + defer sess.Close() + end() + + var results []*CheckResult + for _, c := range checks { + result, err := runCheck(ctx, sess, c, ch.Config.Organization, database, branch) + if err != nil { + // One failing check shouldn't abort the report. Keep the + // catalog next steps so consumers still get the + // cross-reference for checks that timed out or failed. + result = &CheckResult{ + Check: c.Name, + Database: database, + Branch: branch, + Skipped: err.Error(), + NextSteps: formatNextSteps(c.NextSteps, ch.Config.Organization, database, branch), + } + } + results = append(results, result) + + if ch.Printer.Format() == printer.Human { + ch.Printer.Printf("\n%s — %s\n", printer.Bold(c.Name), c.Short) + if result.Skipped != "" { + ch.Printer.Printf(" %s\n", result.Skipped) + printCompactNextSteps(ch, result.NextSteps) + continue + } + printHumanTable(ch, c, result) + // Only surface the targeted follow-up when the check + // actually found something; the combined report already + // ends with the full insights pointer. + if result.RowCount > 0 { + printCompactNextSteps(ch, result.NextSteps) + } + } + } + + report := &Report{ + Database: database, + Branch: branch, + Results: results, + NextSteps: insightsNextSteps(ch.Config.Organization, database, branch), + } + + switch ch.Printer.Format() { + case printer.Human: + ch.Printer.Printf("\nThese checks are point-in-time. For server-side analysis of production traffic\n(slow queries, failing queries, anomalies, schema recommendations), also run:\n") + for _, step := range report.NextSteps { + ch.Printer.Printf(" %s\n", printer.BoldBlue(step)) + } + return nil + case printer.JSON: + return ch.Printer.PrintJSON(report) + default: + return fmt.Errorf("csv output is not supported for inspect all; use a single check or --format json") + } + }, + } +} + +// printNextSteps renders a check's follow-up commands in human output. The +// steps come from formatNextSteps, so they are copy-pasteable (including +// --org and --format json). +func printNextSteps(ch *cmdutil.Helper, steps []string) { + if len(steps) == 0 { + return + } + ch.Printer.Printf("\nFor server-side analysis of production traffic, also see:\n") + for _, step := range steps { + ch.Printer.Printf(" %s\n", printer.BoldBlue(step)) + } +} + +// printCompactNextSteps is the one-line-per-step variant used between +// sections of the combined report. +func printCompactNextSteps(ch *cmdutil.Helper, steps []string) { + for _, step := range steps { + ch.Printer.Printf(" also see: %s\n", printer.BoldBlue(step)) + } +} + +func newSession(ctx context.Context, ch *cmdutil.Helper, database, branch string, flags *inspectFlags) (*sqlquery.Session, error) { + sess, err := sqlquery.NewSession(ctx, ch, sqlquery.Options{ + Organization: ch.Config.Organization, + Database: database, + Branch: branch, + Keyspace: flags.keyspace, + PostgresDB: flags.postgresDB, + PostgresAdditionalRoles: []string{"pg_read_all_stats"}, + Role: flags.role, + Replica: flags.replica, + }) + if err != nil && strings.Contains(err.Error(), "permission denied for database") { + return nil, fmt.Errorf("%w (the reader role may not have CONNECT on database %q; retry with --role admin)", err, flags.postgresDB) + } + return sess, err +} + +func runCheck(ctx context.Context, sess *sqlquery.Session, c check, organization, database, branch string) (*CheckResult, error) { + result := &CheckResult{Check: c.Name, Database: database, Branch: branch} + + var impl *engineSQL + var hint string + switch sess.Engine() { + case "mysql": + impl, hint = c.MySQL, c.MySQLHint + default: + impl, hint = c.Postgres, c.PostgresHint + } + if impl == nil { + if hint == "" { + hint = fmt.Sprintf("The %s check is not available for %s databases.", c.Name, sess.Engine()) + } + result.Skipped = substituteResourceNames(hint, database, branch) + result.NextSteps = formatNextSteps(c.NextSteps, organization, database, branch) + return result, nil + } + + ctx, cancel := context.WithTimeout(ctx, checkTimeout) + defer cancel() + + if impl.RequiresExtension != "" { + // Extension names come from the check catalog above, never from user + // input, so interpolating into the literal is safe. + _, rows, err := sess.Query(ctx, fmt.Sprintf("SELECT 1 FROM pg_extension WHERE extname = '%s'", impl.RequiresExtension)) + if err != nil { + return nil, err + } + if len(rows) == 0 { + msg := fmt.Sprintf("The %s check needs the %q extension, which is not installed.", c.Name, impl.RequiresExtension) + steps := formatNextSteps(c.NextSteps, organization, database, branch) + if len(steps) > 0 { + msg += fmt.Sprintf(" PlanetScale Insights provides this analysis server-side, no extension needed: %s.", strings.Join(steps, "; ")) + } + msg += fmt.Sprintf(" To run this check anyway, enable the extension with: CREATE EXTENSION %s;", impl.RequiresExtension) + result.Skipped = msg + result.NextSteps = steps + return result, nil + } + } + + columns, rows, err := sess.Query(ctx, impl.SQL) + if err != nil { + return nil, fmt.Errorf("%s check failed: %w", c.Name, err) + } + + result.Columns = columns + result.Rows = rows + result.RowCount = len(rows) + result.NextSteps = formatNextSteps(c.NextSteps, organization, database, branch) + return result, nil +} + +func printCheck(ch *cmdutil.Helper, c check, result *CheckResult) error { + switch ch.Printer.Format() { + case printer.JSON: + return ch.Printer.PrintJSON(result) + case printer.CSV: + return printCSV(ch.Printer.ResourceOutput(), result) + default: + printHumanTable(ch, c, result) + return nil + } +} + +func printHumanTable(ch *cmdutil.Helper, c check, result *CheckResult) { + if len(result.Rows) == 0 { + ch.Printer.Printf(" %s\n", c.EmptyMessage) + return + } + + var sb strings.Builder + w := tabwriter.NewWriter(&sb, 2, 2, 2, ' ', 0) + fmt.Fprintln(w, " "+strings.Join(result.Columns, "\t")) + for _, row := range result.Rows { + values := make([]string, 0, len(result.Columns)) + for _, col := range result.Columns { + values = append(values, formatValue(row[col])) + } + fmt.Fprintln(w, " "+strings.Join(values, "\t")) + } + w.Flush() + ch.Printer.Printf("%s", sb.String()) +} + +func printCSV(out io.Writer, result *CheckResult) error { + w := csv.NewWriter(out) + if result.Skipped != "" { + if err := w.Write([]string{"check", "database", "branch", "skipped", "next_steps"}); err != nil { + return err + } + if err := w.Write([]string{ + result.Check, + result.Database, + result.Branch, + formatValue(result.Skipped), + strings.Join(result.NextSteps, "; "), + }); err != nil { + return err + } + } else { + if err := w.Write(result.Columns); err != nil { + return err + } + for _, row := range result.Rows { + values := make([]string, 0, len(result.Columns)) + for _, col := range result.Columns { + values = append(values, formatValue(row[col])) + } + if err := w.Write(values); err != nil { + return err + } + } + } + w.Flush() + return w.Error() +} + +func formatValue(v any) string { + if v == nil { + return "" + } + if s, ok := v.(string); ok { + return strings.Join(strings.Fields(s), " ") + } + return fmt.Sprintf("%v", v) +} diff --git a/internal/cmd/inspect/inspect_output_test.go b/internal/cmd/inspect/inspect_output_test.go new file mode 100644 index 00000000..9696ad3b --- /dev/null +++ b/internal/cmd/inspect/inspect_output_test.go @@ -0,0 +1,39 @@ +package inspect + +import ( + "bytes" + "encoding/csv" + "strings" + "testing" + + qt "github.com/frankban/quicktest" +) + +func TestPrintCSVExplainsSkippedCheck(t *testing.T) { + c := qt.New(t) + + var out bytes.Buffer + err := printCSV(&out, &CheckResult{ + Check: "redundant-indexes", + Database: "mydb", + Branch: "main", + Skipped: "This check is not available for PostgreSQL.", + NextSteps: []string{ + "pscale insights recommendations mydb --org myorg --format json", + }, + }) + c.Assert(err, qt.IsNil) + + records, err := csv.NewReader(strings.NewReader(out.String())).ReadAll() + c.Assert(err, qt.IsNil) + c.Assert(records, qt.DeepEquals, [][]string{ + {"check", "database", "branch", "skipped", "next_steps"}, + { + "redundant-indexes", + "mydb", + "main", + "This check is not available for PostgreSQL.", + "pscale insights recommendations mydb --org myorg --format json", + }, + }) +} diff --git a/internal/cmd/root.go b/internal/cmd/root.go index 8b21e1c5..a56aa24f 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -25,6 +25,7 @@ import ( "strings" "time" + clicontent "github.com/planetscale/cli" "github.com/planetscale/cli/internal/cmd/agentguide" "github.com/planetscale/cli/internal/cmd/mcp" "github.com/planetscale/cli/internal/cmd/role" @@ -43,6 +44,8 @@ import ( "github.com/planetscale/cli/internal/cmd/dataimports" "github.com/planetscale/cli/internal/cmd/deployrequest" "github.com/planetscale/cli/internal/cmd/importcmd" + "github.com/planetscale/cli/internal/cmd/insights" + "github.com/planetscale/cli/internal/cmd/inspect" "github.com/planetscale/cli/internal/cmd/keyspace" "github.com/planetscale/cli/internal/cmd/org" "github.com/planetscale/cli/internal/cmd/password" @@ -232,6 +235,20 @@ func runCmd(ctx context.Context, ver, commit, buildDate string, format *printer. rootCmd.PersistentFlags().Lookup("api-token").DefValue = "" // Add command groups for better organization + // Every command's help/usage output ends with a pointer to the embedded + // agent guide, so agents that discover pscale via --help (rather than a + // repo AGENTS.md) find the machine-readable guidance. + rootCmd.SetUsageTemplate(rootCmd.UsageTemplate() + + "\nAgents: run \"pscale agent-guide --format json\" for machine-readable guidance, or \"pscale help agents\" to read the full guide.\n") + + // Additional help topic (no Run function): `pscale help agents` prints + // the embedded agent guide. + rootCmd.AddCommand(&cobra.Command{ + Use: "agents", + Short: "How AI agents and automation should use pscale", + Long: clicontent.AgentGuide, + }) + rootCmd.AddGroup(&cobra.Group{ID: "database", Title: printer.Bold("MySQL & PostgreSQL database commands:")}) rootCmd.AddGroup(&cobra.Group{ID: "vitess", Title: printer.Bold("Vitess/MySQL-specific commands:")}) rootCmd.AddGroup(&cobra.Group{ID: "postgres", Title: printer.Bold("PostgreSQL-specific commands:")}) @@ -311,6 +328,14 @@ func runCmd(ctx context.Context, ver, commit, buildDate string, format *printer. databaseCmd.GroupID = "database" rootCmd.AddCommand(databaseCmd) + insightsCmd := insights.InsightsCmd(ch) + insightsCmd.GroupID = "database" + rootCmd.AddCommand(insightsCmd) + + inspectCmd := inspect.InspectCmd(ch) + inspectCmd.GroupID = "database" + rootCmd.AddCommand(inspectCmd) + webhookCmd := webhook.WebhookCmd(ch) webhookCmd.GroupID = "database" rootCmd.AddCommand(webhookCmd) diff --git a/internal/cmd/sql/sql.go b/internal/cmd/sql/sql.go index d038940a..91739de1 100644 --- a/internal/cmd/sql/sql.go +++ b/internal/cmd/sql/sql.go @@ -34,7 +34,9 @@ Destructive SQL containing DELETE, DROP, or TRUNCATE is blocked unless --force i Agents must ask the user for approval before using --force. MySQL (Vitess) databases use the primary keyspace by default (same as pscale shell -D @primary). -Pass --keyspace when targeting a specific keyspace in a multi-keyspace database. +Pass --keyspace when targeting a specific keyspace in a multi-keyspace database. A keyspace may +include a shard and tablet type (mykeyspace/-80, mykeyspace/-80@replica) to pin the connection to +one shard; enumerate shards with SHOW VITESS_SHARDS. PostgreSQL databases use --dbname (default postgres). @@ -89,7 +91,7 @@ Place flags after positional arguments (see Usage). --org is required: cmd.PersistentFlags().StringVar(&ch.Config.Organization, "org", ch.Config.Organization, "The organization for the current user") cmd.Flags().StringVar(&flags.query, "query", "", "SQL query to execute") - cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Vitess keyspace (optional; defaults to @primary, same as pscale shell)") + cmd.Flags().StringVar(&flags.keyspace, "keyspace", "", "Vitess keyspace, optionally with a shard and tablet type (e.g. mykeyspace, mykeyspace/-80, mykeyspace/-80@replica). List shards with --query \"SHOW VITESS_SHARDS\". Defaults to @primary, same as pscale shell.") cmd.Flags().StringVar(&flags.postgresDB, "dbname", "postgres", "PostgreSQL database name") cmd.Flags().StringVar(&flags.role, "role", "", "Role defines the access level, allowed values are: reader, writer, readwriter, admin. Defaults to reader (use --role admin for writes).") diff --git a/internal/mock/insights.go b/internal/mock/insights.go new file mode 100644 index 00000000..da668f3a --- /dev/null +++ b/internal/mock/insights.go @@ -0,0 +1,59 @@ +package mock + +import ( + "context" + + ps "github.com/planetscale/cli/internal/planetscale" +) + +type QueryInsightsService struct { + ListQueriesFn func(context.Context, *ps.ListQueryInsightsRequest, ...ps.ListOption) ([]*ps.QueryInsight, error) + ListQueriesFnInvoked bool + + ListErrorsFn func(context.Context, *ps.ListQueryInsightsErrorsRequest, ...ps.ListOption) ([]*ps.QueryInsightError, error) + ListErrorsFnInvoked bool + + ListAnomaliesFn func(context.Context, *ps.ListAnomaliesRequest, ...ps.ListOption) ([]*ps.Anomaly, error) + ListAnomaliesFnInvoked bool +} + +func (s *QueryInsightsService) ListQueries(ctx context.Context, req *ps.ListQueryInsightsRequest, opts ...ps.ListOption) ([]*ps.QueryInsight, error) { + s.ListQueriesFnInvoked = true + return s.ListQueriesFn(ctx, req, opts...) +} + +func (s *QueryInsightsService) ListErrors(ctx context.Context, req *ps.ListQueryInsightsErrorsRequest, opts ...ps.ListOption) ([]*ps.QueryInsightError, error) { + s.ListErrorsFnInvoked = true + return s.ListErrorsFn(ctx, req, opts...) +} + +func (s *QueryInsightsService) ListAnomalies(ctx context.Context, req *ps.ListAnomaliesRequest, opts ...ps.ListOption) ([]*ps.Anomaly, error) { + s.ListAnomaliesFnInvoked = true + return s.ListAnomaliesFn(ctx, req, opts...) +} + +type SchemaRecommendationService struct { + ListFn func(context.Context, *ps.ListSchemaRecommendationsRequest, ...ps.ListOption) ([]*ps.SchemaRecommendation, error) + ListFnInvoked bool + + GetFn func(context.Context, *ps.GetSchemaRecommendationRequest) (*ps.SchemaRecommendation, error) + GetFnInvoked bool + + DismissFn func(context.Context, *ps.DismissSchemaRecommendationRequest) (*ps.SchemaRecommendation, error) + DismissFnInvoked bool +} + +func (s *SchemaRecommendationService) List(ctx context.Context, req *ps.ListSchemaRecommendationsRequest, opts ...ps.ListOption) ([]*ps.SchemaRecommendation, error) { + s.ListFnInvoked = true + return s.ListFn(ctx, req, opts...) +} + +func (s *SchemaRecommendationService) Get(ctx context.Context, req *ps.GetSchemaRecommendationRequest) (*ps.SchemaRecommendation, error) { + s.GetFnInvoked = true + return s.GetFn(ctx, req) +} + +func (s *SchemaRecommendationService) Dismiss(ctx context.Context, req *ps.DismissSchemaRecommendationRequest) (*ps.SchemaRecommendation, error) { + s.DismissFnInvoked = true + return s.DismissFn(ctx, req) +} diff --git a/internal/planetscale/client.go b/internal/planetscale/client.go index be9977db..c40227c2 100644 --- a/internal/planetscale/client.go +++ b/internal/planetscale/client.go @@ -67,6 +67,7 @@ type Client struct { PostgresBranches PostgresBranchesService PostgresRoles PostgresRolesService Processlist ProcesslistService + QueryInsights QueryInsightsService QueryPatterns QueryPatternsService Regions RegionsService SchemaRecommendations SchemaRecommendationService @@ -328,6 +329,7 @@ func NewClient(opts ...ClientOption) (*Client, error) { c.Processlist = &processlistService{client: c} c.PostgresBranches = &postgresBranchesService{client: c} c.PostgresRoles = &postgresRolesService{client: c} + c.QueryInsights = &queryInsightsService{client: c} c.QueryPatterns = &queryPatternsService{client: c} c.Regions = ®ionsService{client: c} c.SchemaRecommendations = &schemaRecommendationService{client: c} diff --git a/internal/planetscale/insights.go b/internal/planetscale/insights.go new file mode 100644 index 00000000..1127443e --- /dev/null +++ b/internal/planetscale/insights.go @@ -0,0 +1,192 @@ +package planetscale + +import ( + "context" + "net/http" + "path" + "time" +) + +var _ QueryInsightsService = &queryInsightsService{} + +// QueryInsightsService is an interface for communicating with the PlanetScale +// Query Insights API: aggregated query statistics, query errors, and detected +// anomalies for a database branch. +type QueryInsightsService interface { + ListQueries(context.Context, *ListQueryInsightsRequest, ...ListOption) ([]*QueryInsight, error) + ListErrors(context.Context, *ListQueryInsightsErrorsRequest, ...ListOption) ([]*QueryInsightError, error) + ListAnomalies(context.Context, *ListAnomaliesRequest, ...ListOption) ([]*Anomaly, error) +} + +// QueryInsight is an aggregated statistics record for a normalized query +// pattern on a branch. +type QueryInsight struct { + ID string `json:"id"` + Fingerprint string `json:"fingerprint"` + NormalizedSQL string `json:"normalized_sql"` + StatementType string `json:"statement_type"` + Keyspace string `json:"keyspace"` + Tables []string `json:"tables"` + IndexUsages []IndexUsage `json:"index_usages"` + QueryCount int64 `json:"query_count"` + ErrorCount int64 `json:"error_count"` + LastRunAt time.Time `json:"last_run_at"` + TimePerQuery float64 `json:"time_per_query"` + P50Latency float64 `json:"p50_latency"` + P99Latency float64 `json:"p99_latency"` + MaxLatency float64 `json:"max_latency"` + SumRowsRead int64 `json:"sum_rows_read"` + SumRowsReturned int64 `json:"sum_rows_returned"` + SumRowsAffected int64 `json:"sum_rows_affected"` + RowsReadPerReturned float64 `json:"rows_read_per_returned"` + SumTotalDurationMillis float64 `json:"sum_total_duration_millis"` + SumTotalDurationPct float64 `json:"sum_total_duration_percent"` + SumCPUDurationMillis float64 `json:"sum_cpu_duration_millis"` + SumIODurationMillis float64 `json:"sum_io_duration_millis"` + BlockCacheHitRatio float64 `json:"block_cache_hit_ratio"` + Multishard bool `json:"multishard"` +} + +// IndexUsage records how often an index served a query pattern. +type IndexUsage struct { + Name string `json:"name"` + Count int64 `json:"count"` + Percent float64 `json:"percent"` +} + +// QueryInsightError is an aggregated error record for a query pattern on a +// branch. +type QueryInsightError struct { + ID string `json:"id"` + ErrorFingerprint string `json:"error_fingerprint"` + StartedAt time.Time `json:"started_at"` + TotalDurationMillis float64 `json:"total_duration_millis"` + TimePerQuery float64 `json:"time_per_query"` + ErrorCount int64 `json:"error_count"` + ErrorMessage string `json:"error_message"` +} + +// Anomaly is a detected resource-usage anomaly on a branch's primary. +type Anomaly struct { + ID string `json:"id"` + PeriodStart time.Time `json:"period_start"` + PeriodEnd time.Time `json:"period_end"` + MinutesInViolation int64 `json:"minutes_in_violation"` + Active bool `json:"active"` + Duration float64 `json:"duration"` + MetricsStart time.Time `json:"metrics_start"` + MetricsEnd time.Time `json:"metrics_end"` +} + +// ListQueryInsightsRequest is the request for listing query statistics for a branch. +type ListQueryInsightsRequest struct { + Organization string + Database string + Branch string +} + +// ListQueryInsightsErrorsRequest is the request for listing query errors for a branch. +type ListQueryInsightsErrorsRequest struct { + Organization string + Database string + Branch string +} + +// ListAnomaliesRequest is the request for listing anomalies for a branch. +type ListAnomaliesRequest struct { + Organization string + Database string + Branch string +} + +// WithSort returns a ListOption that sets the "sort" and "dir" URL parameters. +func WithSort(sort, dir string) ListOption { + return func(opt *ListOptions) error { + if sort != "" { + opt.URLValues.Set("sort", sort) + } + if dir != "" { + opt.URLValues.Set("dir", dir) + } + return nil + } +} + +// WithPeriod returns a ListOption that sets the "period" URL parameter +// (e.g. "1h", "24h"). +func WithPeriod(period string) ListOption { + return func(opt *ListOptions) error { + if period != "" { + opt.URLValues.Set("period", period) + } + return nil + } +} + +type queryInsightsService struct { + client *Client +} + +type queryInsightsResponse struct { + Insights []*QueryInsight `json:"data"` +} + +type queryInsightErrorsResponse struct { + Errors []*QueryInsightError `json:"data"` +} + +type anomaliesResponse struct { + Anomalies []*Anomaly `json:"data"` +} + +func (s *queryInsightsService) ListQueries(ctx context.Context, request *ListQueryInsightsRequest, opts ...ListOption) ([]*QueryInsight, error) { + listOpts := defaultListOptions(opts...) + + req, err := s.client.newRequest(http.MethodGet, insightsAPIPath(request.Organization, request.Database, request.Branch), nil, WithQueryParams(*listOpts.URLValues)) + if err != nil { + return nil, err + } + + resp := &queryInsightsResponse{} + if err := s.client.do(ctx, req, &resp); err != nil { + return nil, err + } + + return resp.Insights, nil +} + +func (s *queryInsightsService) ListErrors(ctx context.Context, request *ListQueryInsightsErrorsRequest, opts ...ListOption) ([]*QueryInsightError, error) { + listOpts := defaultListOptions(opts...) + + req, err := s.client.newRequest(http.MethodGet, path.Join(insightsAPIPath(request.Organization, request.Database, request.Branch), "errors"), nil, WithQueryParams(*listOpts.URLValues)) + if err != nil { + return nil, err + } + + resp := &queryInsightErrorsResponse{} + if err := s.client.do(ctx, req, &resp); err != nil { + return nil, err + } + + return resp.Errors, nil +} + +func (s *queryInsightsService) ListAnomalies(ctx context.Context, request *ListAnomaliesRequest, opts ...ListOption) ([]*Anomaly, error) { + listOpts := defaultListOptions(opts...) + + req, err := s.client.newRequest(http.MethodGet, path.Join(insightsAPIPath(request.Organization, request.Database, request.Branch), "anomalies"), nil, WithQueryParams(*listOpts.URLValues)) + if err != nil { + return nil, err + } + + resp := &anomaliesResponse{} + if err := s.client.do(ctx, req, &resp); err != nil { + return nil, err + } + + return resp.Anomalies, nil +} + +func insightsAPIPath(org, db, branch string) string { + return path.Join("v1/organizations", org, "databases", db, "branches", branch, "insights") +} diff --git a/internal/planetscale/insights_test.go b/internal/planetscale/insights_test.go new file mode 100644 index 00000000..93e9e41a --- /dev/null +++ b/internal/planetscale/insights_test.go @@ -0,0 +1,168 @@ +package planetscale + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + qt "github.com/frankban/quicktest" +) + +func TestQueryInsights_ListQueries(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + c.Assert(r.Method, qt.Equals, http.MethodGet) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/planetscale-go-test-db/branches/main/insights") + c.Assert(r.URL.Query().Get("sort"), qt.Equals, "totalTime") + c.Assert(r.URL.Query().Get("dir"), qt.Equals, "desc") + c.Assert(r.URL.Query().Get("per_page"), qt.Equals, "5") + c.Assert(r.URL.Query().Get("period"), qt.Equals, "24h") + + out := `{ + "type": "list", + "current_page": 1, + "data": [{ + "id": "f5a67b04ee9f", + "type": "HourlyBranchQuery", + "query_count": 8, + "error_count": 1, + "tables": ["users"], + "index_usages": [{"name": "public.users.users_pkey", "count": 5757, "percent": 100.0}], + "sum_rows_read": 100, + "sum_rows_returned": 25, + "sum_rows_affected": 0, + "rows_read_per_returned": 4.0, + "sum_total_duration_millis": 14.595, + "sum_total_duration_percent": 57.57, + "sum_cpu_duration_millis": 1.5, + "sum_io_duration_millis": 0.5, + "last_run_at": "2026-07-22T23:22:14.000Z", + "time_per_query": 1.824375, + "p50_latency": 1.75, + "p99_latency": 2.32, + "max_latency": 2.32, + "block_cache_hit_ratio": 0.9, + "fingerprint": "b129e8fa", + "statement_type": "SELECT", + "keyspace": "mydb", + "normalized_sql": "select * from users where id = ?", + "multishard": false + }] + }` + _, err := w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + insights, err := client.QueryInsights.ListQueries(context.Background(), &ListQueryInsightsRequest{ + Organization: testOrg, + Database: testDatabase, + Branch: "main", + }, WithPerPage(5), WithSort("totalTime", "desc"), WithPeriod("24h")) + + c.Assert(err, qt.IsNil) + c.Assert(len(insights), qt.Equals, 1) + c.Assert(insights[0].ID, qt.Equals, "f5a67b04ee9f") + c.Assert(insights[0].QueryCount, qt.Equals, int64(8)) + c.Assert(insights[0].ErrorCount, qt.Equals, int64(1)) + c.Assert(insights[0].NormalizedSQL, qt.Equals, "select * from users where id = ?") + c.Assert(insights[0].StatementType, qt.Equals, "SELECT") + c.Assert(insights[0].Keyspace, qt.Equals, "mydb") + c.Assert(insights[0].SumRowsRead, qt.Equals, int64(100)) + c.Assert(insights[0].RowsReadPerReturned, qt.Equals, 4.0) + c.Assert(insights[0].SumTotalDurationMillis, qt.Equals, 14.595) + c.Assert(insights[0].P99Latency, qt.Equals, 2.32) + c.Assert(insights[0].Tables, qt.DeepEquals, []string{"users"}) + c.Assert(insights[0].IndexUsages, qt.DeepEquals, []IndexUsage{{Name: "public.users.users_pkey", Count: 5757, Percent: 100.0}}) +} + +func TestQueryInsights_ListErrors(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + c.Assert(r.Method, qt.Equals, http.MethodGet) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/planetscale-go-test-db/branches/main/insights/errors") + + out := `{ + "type": "list", + "current_page": 1, + "data": [{ + "id": "5d4485f99812", + "type": "BranchQuery", + "error_fingerprint": "5d4485f9981294d1", + "started_at": "2026-07-22T14:05:50.000Z", + "total_duration_millis": 12.5, + "time_per_query": 3.1, + "error_count": 4, + "error_message": "relation \"widgets\" does not exist" + }] + }` + _, err := w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + errs, err := client.QueryInsights.ListErrors(context.Background(), &ListQueryInsightsErrorsRequest{ + Organization: testOrg, + Database: testDatabase, + Branch: "main", + }) + + c.Assert(err, qt.IsNil) + c.Assert(len(errs), qt.Equals, 1) + c.Assert(errs[0].ID, qt.Equals, "5d4485f99812") + c.Assert(errs[0].ErrorCount, qt.Equals, int64(4)) + c.Assert(errs[0].ErrorMessage, qt.Equals, `relation "widgets" does not exist`) + c.Assert(errs[0].TotalDurationMillis, qt.Equals, 12.5) +} + +func TestQueryInsights_ListAnomalies(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + c.Assert(r.Method, qt.Equals, http.MethodGet) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/planetscale-go-test-db/branches/main/insights/anomalies") + + out := `{ + "type": "list", + "current_page": 1, + "data": [{ + "id": "anomaly-123", + "period_start": "2026-07-22T14:00:00.000Z", + "period_end": "2026-07-22T14:30:00.000Z", + "minutes_in_violation": 12, + "active": false, + "duration": 1800.0, + "metrics_start": "2026-07-22T13:30:00.000Z", + "metrics_end": "2026-07-22T15:00:00.000Z" + }] + }` + _, err := w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + })) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + anomalies, err := client.QueryInsights.ListAnomalies(context.Background(), &ListAnomaliesRequest{ + Organization: testOrg, + Database: testDatabase, + Branch: "main", + }) + + c.Assert(err, qt.IsNil) + c.Assert(len(anomalies), qt.Equals, 1) + c.Assert(anomalies[0].ID, qt.Equals, "anomaly-123") + c.Assert(anomalies[0].MinutesInViolation, qt.Equals, int64(12)) + c.Assert(anomalies[0].Active, qt.Equals, false) + c.Assert(anomalies[0].Duration, qt.Equals, 1800.0) +} diff --git a/internal/printer/printer.go b/internal/printer/printer.go index b5c35d67..cd4da5e2 100644 --- a/internal/printer/printer.go +++ b/internal/printer/printer.go @@ -192,6 +192,16 @@ func (p *Printer) SetResourceOutput(out io.Writer) { p.resourceOut = out } +// ResourceOutput returns the writer PrintResource writes to, for commands that +// render resources themselves (e.g. dynamic column sets that can't use struct +// tags). +func (p *Printer) ResourceOutput() io.Writer { + if p.resourceOut != nil { + return p.resourceOut + } + return os.Stdout +} + // PrintResource prints the given resource in the format it was specified. func (p *Printer) PrintResource(v interface{}) error { if p.format == nil { diff --git a/internal/sqlquery/dsn_test.go b/internal/sqlquery/dsn_test.go new file mode 100644 index 00000000..9dfabfc6 --- /dev/null +++ b/internal/sqlquery/dsn_test.go @@ -0,0 +1,27 @@ +package sqlquery + +import ( + "testing" + + qt "github.com/frankban/quicktest" + gomysql "github.com/go-sql-driver/mysql" +) + +// Shard-targeted keyspace names contain characters ("/", "@", "-") that must +// survive the round trip through the driver DSN. +func TestMySQLDSNEscapesShardTargets(t *testing.T) { + c := qt.New(t) + + for _, dbName := range []string{ + "@primary", + "mykeyspace", + "mykeyspace/-80", + "mykeyspace/80-c0@replica", + "mykeyspace/-80@primary", + } { + dsn := mysqlDSN("127.0.0.1:3306", Options{Keyspace: dbName}) + parsed, err := gomysql.ParseDSN(dsn) + c.Assert(err, qt.IsNil, qt.Commentf("dbname %q", dbName)) + c.Assert(parsed.DBName, qt.Equals, dbName) + } +} diff --git a/internal/sqlquery/session.go b/internal/sqlquery/session.go new file mode 100644 index 00000000..017074c7 --- /dev/null +++ b/internal/sqlquery/session.go @@ -0,0 +1,113 @@ +package sqlquery + +import ( + "context" + "database/sql" + "fmt" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" +) + +// Session is an open connection to a database branch using ephemeral +// credentials. Unlike Execute, which opens and tears down a connection per +// query, a Session lets callers run several queries (e.g. a suite of +// diagnostic checks) over a single connection. Close must be called to release +// the connection and clean up the ephemeral credentials. +type Session struct { + // Kind is the database kind as reported by the API: "mysql", + // "postgresql", or "horizon". + Kind string + + db *sql.DB + cleanup func() +} + +// NewSession validates the options, looks up the database and branch, and +// opens a connection with ephemeral credentials (a branch password for MySQL, +// a temporary role for PostgreSQL). +func NewSession(ctx context.Context, ch *cmdutil.Helper, opts Options) (*Session, error) { + if opts.Organization == "" { + return nil, fmt.Errorf("organization is required (use --org or set org in pscale.yml)") + } + if opts.Database == "" || opts.Branch == "" { + return nil, fmt.Errorf("database and branch are required") + } + + role, err := cmdutil.ResolveAccessRole(opts.Role, opts.Replica, cmdutil.ReaderRole) + if err != nil { + return nil, err + } + + client, err := ch.Client() + if err != nil { + return nil, err + } + + dbInfo, err := client.Databases.Get(ctx, &ps.GetDatabaseRequest{ + Organization: opts.Organization, + Database: opts.Database, + }) + if err != nil { + return nil, fmt.Errorf("database lookup: %w", err) + } + + dbBranch, err := client.DatabaseBranches.Get(ctx, &ps.GetDatabaseBranchRequest{ + Organization: opts.Organization, + Database: opts.Database, + Branch: opts.Branch, + }) + if err != nil { + return nil, fmt.Errorf("branch lookup: %w", err) + } + if !dbBranch.Ready { + return nil, fmt.Errorf("database branch is not ready yet") + } + + var db *sql.DB + var cleanup func() + + switch string(dbInfo.Kind) { + case "mysql": + db, cleanup, err = openMySQL(ctx, ch, opts, role) + case "postgresql", "horizon": + pgDB := opts.PostgresDB + if pgDB == "" { + pgDB = "postgres" + } + db, cleanup, err = openPostgres(ctx, ch, opts, pgDB, role) + default: + return nil, fmt.Errorf("unsupported database kind %q", dbInfo.Kind) + } + if err != nil { + return nil, err + } + + return &Session{Kind: string(dbInfo.Kind), db: db, cleanup: cleanup}, nil +} + +// Engine normalizes Kind to "mysql" or "postgresql". +func (s *Session) Engine() string { + if s.Kind == "mysql" { + return "mysql" + } + return "postgresql" +} + +// Query runs a single query over the session's connection and returns the +// column names (in result order) and rows. +func (s *Session) Query(ctx context.Context, query string) ([]string, []map[string]any, error) { + outcome, err := runQuery(ctx, s.db, query) + if err != nil { + return nil, nil, err + } + return outcome.columns, outcome.rows, nil +} + +// Close releases the connection and cleans up the ephemeral credentials. +func (s *Session) Close() { + if s.cleanup != nil { + s.cleanup() + s.cleanup = nil + } +} diff --git a/internal/sqlquery/sqlquery.go b/internal/sqlquery/sqlquery.go index df7143b9..e935bb65 100644 --- a/internal/sqlquery/sqlquery.go +++ b/internal/sqlquery/sqlquery.go @@ -9,7 +9,7 @@ import ( "strings" "time" - _ "github.com/go-sql-driver/mysql" + gomysql "github.com/go-sql-driver/mysql" _ "github.com/lib/pq" "github.com/planetscale/cli/internal/cmdutil" @@ -30,6 +30,8 @@ type Options struct { Keyspace string // PostgresDB is the PostgreSQL database name. Defaults to postgres. PostgresDB string + // PostgresAdditionalRoles adds built-in roles to the ephemeral credential. + PostgresAdditionalRoles []string // Role is reader, writer, readwriter, or admin (same as pscale shell --role). Role string // Replica routes reads to replicas when true (same as pscale shell --replica). @@ -142,10 +144,23 @@ func Execute(ctx context.Context, ch *cmdutil.Helper, opts Options) (*Result, er } func queryMySQL(ctx context.Context, ch *cmdutil.Helper, opts Options, role cmdutil.PasswordRole) (*queryOutcome, error) { - client, err := ch.Client() + db, cleanup, err := openMySQL(ctx, ch, opts, role) if err != nil { return nil, err } + defer cleanup() + + return runQuery(ctx, db, opts.Query) +} + +// openMySQL mints an ephemeral branch password, starts an in-process proxy, +// and opens a connection through it. The returned cleanup closes the +// connection and proxy and deletes the password. +func openMySQL(ctx context.Context, ch *cmdutil.Helper, opts Options, role cmdutil.PasswordRole) (*sql.DB, func(), error) { + client, err := ch.Client() + if err != nil { + return nil, nil, err + } pw, err := passwordutil.New(ctx, client, passwordutil.Options{ Organization: opts.Organization, @@ -157,13 +172,13 @@ func queryMySQL(ctx context.Context, ch *cmdutil.Helper, opts Options, role cmdu Replica: opts.Replica, }) if err != nil { - return nil, err + return nil, nil, err } - defer func() { + cleanupPassword := func() { cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = pw.Cleanup(cleanupCtx) - }() + } proxy := proxyutil.New(proxyutil.Config{ Logger: cmdutil.NewZapLogger(ch.Debug()), @@ -171,36 +186,49 @@ func queryMySQL(ctx context.Context, ch *cmdutil.Helper, opts Options, role cmdu Username: pw.Password.Username, Password: pw.Password.PlainText, }) - defer proxy.Close() l, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { - return nil, err + proxy.Close() + cleanupPassword() + return nil, nil, err } - defer l.Close() errCh := make(chan error, 1) go func() { errCh <- proxy.Serve(l, mysql.CachingSha2Password) }() - db, err := sql.Open("mysql", fmt.Sprintf("root@tcp(%s)/%s", l.Addr().String(), mysqlDSNDatabase(opts))) + // Config.FormatDSN escapes shard targets such as keyspace/-80@replica. + dsn := mysqlDSN(l.Addr().String(), opts) + + db, err := sql.Open("mysql", dsn) if err != nil { - return nil, err + proxy.Close() + l.Close() + <-errCh + cleanupPassword() + return nil, nil, err } - defer db.Close() db.SetConnMaxLifetime(30 * time.Second) - outcome, err := runQuery(ctx, db, opts.Query) - if err != nil { - return nil, err + cleanup := func() { + db.Close() + proxy.Close() + l.Close() + <-errCh + cleanupPassword() } + return db, cleanup, nil +} - proxy.Close() - l.Close() - <-errCh - - return outcome, nil +func mysqlDSN(addr string, opts Options) string { + dsnCfg := gomysql.NewConfig() + dsnCfg.User = "root" + dsnCfg.Net = "tcp" + dsnCfg.Addr = addr + dsnCfg.DBName = mysqlDSNDatabase(opts) + return dsnCfg.FormatDSN() } // mysqlDSNDatabase picks the MySQL database name in the DSN, matching pscale shell: @@ -216,12 +244,25 @@ func mysqlDSNDatabase(opts Options) string { } func queryPostgres(ctx context.Context, ch *cmdutil.Helper, opts Options, pgDB string, role cmdutil.PasswordRole) (*queryOutcome, error) { - client, err := ch.Client() + db, cleanup, err := openPostgres(ctx, ch, opts, pgDB, role) if err != nil { return nil, err } + defer cleanup() + + return runQuery(ctx, db, opts.Query) +} + +// openPostgres mints an ephemeral role and opens a direct connection to the +// branch. The returned cleanup closes the connection and deletes the role. +func openPostgres(ctx context.Context, ch *cmdutil.Helper, opts Options, pgDB string, role cmdutil.PasswordRole) (*sql.DB, func(), error) { + client, err := ch.Client() + if err != nil { + return nil, nil, err + } inheritedRoles, successor := cmdutil.PostgresInheritedRoles(role) + inheritedRoles = append(inheritedRoles, opts.PostgresAdditionalRoles...) pgRole, err := roleutil.New(ctx, client, roleutil.Options{ Organization: opts.Organization, @@ -232,13 +273,13 @@ func queryPostgres(ctx context.Context, ch *cmdutil.Helper, opts Options, pgDB s InheritedRoles: inheritedRoles, }) if err != nil { - return nil, err + return nil, nil, err } - defer func() { + cleanupRole := func() { cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = pgRole.Cleanup(cleanupCtx, successor) - }() + } username := pgRole.Role.Username if opts.Replica { @@ -256,16 +297,22 @@ func queryPostgres(ctx context.Context, ch *cmdutil.Helper, opts Options, pgDB s db, err := sql.Open("postgres", connStr) if err != nil { - return nil, err + cleanupRole() + return nil, nil, err } - defer db.Close() db.SetConnMaxLifetime(30 * time.Second) if err := db.PingContext(ctx); err != nil { - return nil, err + db.Close() + cleanupRole() + return nil, nil, err } - return runQuery(ctx, db, opts.Query) + cleanup := func() { + db.Close() + cleanupRole() + } + return db, cleanup, nil } var readQueryPrefixes = []string{"SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "TABLE"}