From 4765545a6ac53318ae2a5424492b092da905ca7b Mon Sep 17 00:00:00 2001 From: bwarminski Date: Wed, 29 Jul 2026 14:59:21 -0400 Subject: [PATCH 1/3] Add auth attempt export API client --- internal/planetscale/auth_attempt_exports.go | 143 ++++++++++++++ .../planetscale/auth_attempt_exports_test.go | 181 ++++++++++++++++++ internal/planetscale/client.go | 14 +- internal/planetscale/query_patterns.go | 55 +----- internal/planetscale/signed_download.go | 52 +++++ 5 files changed, 395 insertions(+), 50 deletions(-) create mode 100644 internal/planetscale/auth_attempt_exports.go create mode 100644 internal/planetscale/auth_attempt_exports_test.go create mode 100644 internal/planetscale/signed_download.go diff --git a/internal/planetscale/auth_attempt_exports.go b/internal/planetscale/auth_attempt_exports.go new file mode 100644 index 00000000..a2266289 --- /dev/null +++ b/internal/planetscale/auth_attempt_exports.go @@ -0,0 +1,143 @@ +package planetscale + +import ( + "context" + "fmt" + "io" + "net/http" + "path" + "strconv" + "strings" + "time" +) + +type AuthAttemptExportFilters struct { + SourceIPs []string `json:"source_ips,omitempty"` + Branches []string `json:"branches,omitempty"` + Outcomes []string `json:"outcomes,omitempty"` + Usernames []string `json:"usernames,omitempty"` + StartupDatabases []string `json:"startup_databases,omitempty"` + FailureReasons []string `json:"failure_reasons,omitempty"` + BackendRoutes []string `json:"backend_routes,omitempty"` +} + +type CreateAuthAttemptExportRequest struct { + Organization string `json:"-"` + StartAt time.Time `json:"start_at"` + EndAt time.Time `json:"end_at"` + Format string `json:"format"` + Filters AuthAttemptExportFilters `json:"filters,omitempty"` +} + +type GetAuthAttemptExportRequest struct { + Organization string + Export string +} + +type DownloadAuthAttemptExportRequest struct { + Organization string + Export string +} + +type AuthAttemptExport struct { + PublicID string `json:"id"` + State string `json:"state"` + StartAt time.Time `json:"start_at"` + EndAt time.Time `json:"end_at"` + Filters AuthAttemptExportFilters `json:"filters"` + Format string `json:"format"` + ResolvedBranchPublicIDs []string `json:"resolved_branch_public_ids"` + CreatedAt time.Time `json:"created_at"` + StartedAt *time.Time `json:"started_at"` + GeneratedAt *time.Time `json:"generated_at"` + FinishedAt *time.Time `json:"finished_at"` + ExpiresAt *time.Time `json:"expires_at"` + FailureReason string `json:"failure_reason"` + FailureDetail string `json:"failure_detail"` + RecoveryHint string `json:"recovery_hint"` + RetryAfter time.Duration `json:"-"` +} + +type AuthAttemptExportsService interface { + CreateExport(context.Context, *CreateAuthAttemptExportRequest) (*AuthAttemptExport, error) + GetExport(context.Context, *GetAuthAttemptExportRequest) (*AuthAttemptExport, error) + DownloadExport(context.Context, *DownloadAuthAttemptExportRequest) (io.ReadCloser, error) +} + +type authAttemptExportsService struct { + client *Client +} + +var _ AuthAttemptExportsService = &authAttemptExportsService{} + +func (s *authAttemptExportsService) CreateExport(ctx context.Context, createReq *CreateAuthAttemptExportRequest) (*AuthAttemptExport, error) { + req, err := s.client.newRequest(http.MethodPost, authAttemptExportsAPIPath(createReq.Organization), createReq) + if err != nil { + return nil, fmt.Errorf("error creating http request: %w", err) + } + + export := &AuthAttemptExport{} + headers, err := s.client.doWithHeaders(ctx, req, export) + if err != nil { + return nil, err + } + export.RetryAfter = parseRetryAfter(headers.Get("Retry-After")) + return export, nil +} + +func (s *authAttemptExportsService) GetExport(ctx context.Context, getReq *GetAuthAttemptExportRequest) (*AuthAttemptExport, error) { + req, err := s.client.newRequest(http.MethodGet, authAttemptExportAPIPath(getReq.Organization, getReq.Export), nil) + if err != nil { + return nil, fmt.Errorf("error creating http request: %w", err) + } + + export := &AuthAttemptExport{} + headers, err := s.client.doWithHeaders(ctx, req, export) + if err != nil { + return nil, err + } + export.RetryAfter = parseRetryAfter(headers.Get("Retry-After")) + return export, nil +} + +func (s *authAttemptExportsService) DownloadExport(ctx context.Context, downloadReq *DownloadAuthAttemptExportRequest) (io.ReadCloser, error) { + reqPath := path.Join(authAttemptExportAPIPath(downloadReq.Organization, downloadReq.Export), "download") + req, err := s.client.newRequest(http.MethodGet, reqPath, nil) + if err != nil { + return nil, fmt.Errorf("error creating http request: %w", err) + } + body, err := s.client.downloadSignedURL(ctx, req) + if err != nil { + return nil, fmt.Errorf("downloading auth attempt export: %w", err) + } + return body, nil +} + +func authAttemptExportsAPIPath(org string) string { + return path.Join("v1/organizations", org, "auth-attempt-exports") +} + +func authAttemptExportAPIPath(org, export string) string { + return path.Join(authAttemptExportsAPIPath(org), export) +} + +func parseRetryAfter(value string) time.Duration { + value = strings.TrimSpace(value) + if value == "" { + return 0 + } + + seconds, err := strconv.ParseUint(value, 10, 64) + if err == nil { + if seconds > 0 && seconds <= uint64((time.Duration(1<<63-1))/time.Second) { + return time.Duration(seconds) * time.Second + } + return 0 + } + + when, err := http.ParseTime(value) + if err != nil || !when.After(time.Now()) { + return 0 + } + return time.Until(when) +} diff --git a/internal/planetscale/auth_attempt_exports_test.go b/internal/planetscale/auth_attempt_exports_test.go new file mode 100644 index 00000000..0cf10ba0 --- /dev/null +++ b/internal/planetscale/auth_attempt_exports_test.go @@ -0,0 +1,181 @@ +package planetscale + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + qt "github.com/frankban/quicktest" +) + +func TestAuthAttemptExports_CreateExport(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.Method, qt.Equals, http.MethodPost) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/auth-attempt-exports") + + var got map[string]any + c.Assert(json.NewDecoder(r.Body).Decode(&got), qt.IsNil) + want := map[string]any{ + "start_at": "2026-07-29T00:00:00Z", + "end_at": "2026-07-29T01:00:00Z", + "format": "parquet", + "filters": map[string]any{ + "source_ips": []any{"203.0.113.0/24", "2001:db8::/112"}, + "branches": []any{"db/production"}, + "outcomes": []any{"deny"}, + "usernames": []any{"incident-user"}, + "startup_databases": []any{""}, + "failure_reasons": []any{"bad_password"}, + "backend_routes": []any{"postgres"}, + }, + } + c.Assert(got, qt.DeepEquals, want) + + w.Header().Set("Retry-After", "5") + w.WriteHeader(http.StatusCreated) + _, err := w.Write([]byte(`{"id":"export1","state":"pending","format":"parquet"}`)) + c.Assert(err, qt.IsNil) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + export, err := client.AuthAttemptExports.CreateExport(context.Background(), &CreateAuthAttemptExportRequest{ + Organization: "my-org", + StartAt: time.Date(2026, time.July, 29, 0, 0, 0, 0, time.UTC), + EndAt: time.Date(2026, time.July, 29, 1, 0, 0, 0, time.UTC), + Format: "parquet", + Filters: AuthAttemptExportFilters{ + SourceIPs: []string{"203.0.113.0/24", "2001:db8::/112"}, + Branches: []string{"db/production"}, + Outcomes: []string{"deny"}, + Usernames: []string{"incident-user"}, + StartupDatabases: []string{""}, + FailureReasons: []string{"bad_password"}, + BackendRoutes: []string{"postgres"}, + }, + }) + + c.Assert(err, qt.IsNil) + c.Assert(export.PublicID, qt.Equals, "export1") + c.Assert(export.State, qt.Equals, "pending") + c.Assert(export.Format, qt.Equals, "parquet") + c.Assert(export.RetryAfter, qt.Equals, 5*time.Second) +} + +func TestAuthAttemptExports_GetExport(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.Method, qt.Equals, http.MethodGet) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/auth-attempt-exports/export1") + + w.Header().Set("Retry-After", "2") + w.WriteHeader(http.StatusOK) + _, err := w.Write([]byte(`{ + "id":"export1", + "state":"failed", + "start_at":"2026-07-29T00:00:00Z", + "end_at":"2026-07-29T01:00:00Z", + "format":"parquet", + "filters":{ + "source_ips":["203.0.113.0/24"], + "branches":["db/production"], + "outcomes":["deny"], + "usernames":["incident-user"], + "startup_databases":[""], + "failure_reasons":["bad_password"], + "backend_routes":["postgres"] + }, + "resolved_branch_public_ids":["branch1"], + "created_at":"2026-07-29T01:01:00Z", + "started_at":"2026-07-29T01:02:00Z", + "generated_at":"2026-07-29T01:03:00Z", + "finished_at":"2026-07-29T01:04:00Z", + "expires_at":"2026-07-30T01:04:00Z", + "failure_reason":"generation_failed", + "failure_detail":"The export could not be generated.", + "recovery_hint":"Re-create the export; if the error repeats, contact support." + }`)) + c.Assert(err, qt.IsNil) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + export, err := client.AuthAttemptExports.GetExport(context.Background(), &GetAuthAttemptExportRequest{ + Organization: "my-org", + Export: "export1", + }) + + startedAt := time.Date(2026, time.July, 29, 1, 2, 0, 0, time.UTC) + generatedAt := time.Date(2026, time.July, 29, 1, 3, 0, 0, time.UTC) + finishedAt := time.Date(2026, time.July, 29, 1, 4, 0, 0, time.UTC) + expiresAt := time.Date(2026, time.July, 30, 1, 4, 0, 0, time.UTC) + want := &AuthAttemptExport{ + PublicID: "export1", + State: "failed", + StartAt: time.Date(2026, time.July, 29, 0, 0, 0, 0, time.UTC), + EndAt: time.Date(2026, time.July, 29, 1, 0, 0, 0, time.UTC), + Format: "parquet", + Filters: AuthAttemptExportFilters{ + SourceIPs: []string{"203.0.113.0/24"}, + Branches: []string{"db/production"}, + Outcomes: []string{"deny"}, + Usernames: []string{"incident-user"}, + StartupDatabases: []string{""}, + FailureReasons: []string{"bad_password"}, + BackendRoutes: []string{"postgres"}, + }, + ResolvedBranchPublicIDs: []string{"branch1"}, + CreatedAt: time.Date(2026, time.July, 29, 1, 1, 0, 0, time.UTC), + StartedAt: &startedAt, + GeneratedAt: &generatedAt, + FinishedAt: &finishedAt, + ExpiresAt: &expiresAt, + FailureReason: "generation_failed", + FailureDetail: "The export could not be generated.", + RecoveryHint: "Re-create the export; if the error repeats, contact support.", + RetryAfter: 2 * time.Second, + } + + c.Assert(err, qt.IsNil) + c.Assert(export, qt.DeepEquals, want) +} + +func TestParseRetryAfter(t *testing.T) { + c := qt.New(t) + future := time.Now().Add(2 * time.Minute).UTC().Format(http.TimeFormat) + past := time.Now().Add(-time.Minute).UTC().Format(http.TimeFormat) + tests := []struct { + name string + header string + want time.Duration + }{ + {name: "positive delta seconds", header: "5", want: 5 * time.Second}, + {name: "future HTTP date", header: future}, + {name: "absent", header: "", want: 0}, + {name: "malformed", header: "soon", want: 0}, + {name: "fractional", header: "1.5", want: 0}, + {name: "zero", header: "0", want: 0}, + {name: "past HTTP date", header: past, want: 0}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got := parseRetryAfter(test.header) + if test.name == "future HTTP date" { + c.Assert(got > 0, qt.IsTrue) + return + } + c.Assert(got, qt.Equals, test.want) + }) + } +} diff --git a/internal/planetscale/client.go b/internal/planetscale/client.go index ae2f4303..05926807 100644 --- a/internal/planetscale/client.go +++ b/internal/planetscale/client.go @@ -50,6 +50,7 @@ type Client struct { baseURL *url.URL AuditLogs AuditLogsService + AuthAttemptExports AuthAttemptExportsService Backups BackupsService BranchInfrastructure BranchInfrastructureService D1ImportNotifications D1ImportNotificationsService @@ -318,6 +319,7 @@ func NewClient(opts ...ClientOption) (*Client, error) { c.client.CheckRedirect = makeSameHostCheckRedirect(c.baseURL.Hostname()) c.AuditLogs = &auditlogsService{client: c} + c.AuthAttemptExports = &authAttemptExportsService{client: c} c.Backups = &backupsService{client: c} c.BranchInfrastructure = &branchInfrastructureService{client: c} c.D1ImportNotifications = &d1ImportNotificationsService{client: c} @@ -352,14 +354,18 @@ func NewClient(opts ...ClientOption) (*Client, error) { // do makes an HTTP request and populates the given struct v from the response. func (c *Client) do(ctx context.Context, req *http.Request, v interface{}) error { - req = req.WithContext(ctx) - res, err := c.client.Do(req) + _, err := c.doWithHeaders(ctx, req, v) + return err +} + +func (c *Client) doWithHeaders(ctx context.Context, req *http.Request, v interface{}) (http.Header, error) { + res, err := c.client.Do(req.WithContext(ctx)) if err != nil { - return err + return nil, err } defer res.Body.Close() - return c.handleResponse(ctx, res, v) + return res.Header, c.handleResponse(ctx, res, v) } // handleResponse makes an HTTP request and populates the given struct v from diff --git a/internal/planetscale/query_patterns.go b/internal/planetscale/query_patterns.go index 20b669d5..93c0937e 100644 --- a/internal/planetscale/query_patterns.go +++ b/internal/planetscale/query_patterns.go @@ -10,8 +10,6 @@ import ( "net/http" "path" "time" - - "github.com/hashicorp/go-cleanhttp" ) // QueryPatternsReport represents a query patterns report for a branch. @@ -74,7 +72,7 @@ func (s *queryPatternsService) CreateReport(ctx context.Context, createReq *Crea } report := &QueryPatternsReport{} - if err := s.client.do(ctx, req, &report); err != nil { + if err := s.client.do(ctx, req, report); err != nil { return nil, err } @@ -90,7 +88,7 @@ func (s *queryPatternsService) GetReport(ctx context.Context, getReq *GetQueryPa } report := &QueryPatternsReport{} - if err := s.client.do(ctx, req, &report); err != nil { + if err := s.client.do(ctx, req, report); err != nil { return nil, err } @@ -106,49 +104,11 @@ func (s *queryPatternsService) DownloadReport(ctx context.Context, downloadReq * return nil, fmt.Errorf("error creating http request: %w", err) } - // The download endpoint redirects to blob storage. The client's - // credentials live in its transport, so following the redirect with that - // client would send the Authorization header to the storage host, which - // rejects requests carrying credentials beyond the presigned URL. Stop at - // the redirect and fetch its target with an unauthenticated client. - httpClient := *s.client.client - httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { - return http.ErrUseLastResponse - } - - res, err := httpClient.Do(req.WithContext(ctx)) + body, err := s.client.downloadSignedURL(ctx, req) if err != nil { - return nil, err - } - - switch { - case res.StatusCode >= 300 && res.StatusCode < 400: - location, err := res.Location() - res.Body.Close() - if err != nil { - return nil, err - } - - blobReq, err := http.NewRequestWithContext(ctx, http.MethodGet, location.String(), nil) - if err != nil { - return nil, err - } - blobReq.Header.Set("User-Agent", s.client.UserAgent) - - res, err = cleanhttp.DefaultClient().Do(blobReq) - if err != nil { - return nil, err - } - if res.StatusCode >= 300 { - res.Body.Close() - return nil, fmt.Errorf("downloading query patterns report: %s", http.StatusText(res.StatusCode)) - } - case res.StatusCode >= 400: - defer res.Body.Close() - return nil, s.client.handleResponse(ctx, res, nil) + return nil, fmt.Errorf("downloading query patterns report: %w", err) } - - return decompressedReadCloser(res.Body) + return decompressedReadCloser(body) } // decompressedReadCloser wraps body so gzip-compressed content is transparently @@ -170,6 +130,7 @@ func decompressedReadCloser(body io.ReadCloser) (io.ReadCloser, error) { return nil, err } content = gz + return &bodyReadCloser{Reader: content, body: body}, nil } return &bodyReadCloser{Reader: content, body: body}, nil @@ -180,7 +141,9 @@ type bodyReadCloser struct { body io.ReadCloser } -func (r *bodyReadCloser) Close() error { return r.body.Close() } +func (r *bodyReadCloser) Close() error { + return r.body.Close() +} func queryPatternsAPIPath(org, db, branch string) string { return path.Join(databaseBranchAPIPath(org, db, branch), "query-patterns") diff --git a/internal/planetscale/signed_download.go b/internal/planetscale/signed_download.go new file mode 100644 index 00000000..e41a22ff --- /dev/null +++ b/internal/planetscale/signed_download.go @@ -0,0 +1,52 @@ +package planetscale + +import ( + "context" + "fmt" + "io" + "net/http" + + "github.com/hashicorp/go-cleanhttp" +) + +// downloadSignedURL uses an anonymous client to keep API credentials isolated to the request using them +func (c *Client) downloadSignedURL(ctx context.Context, req *http.Request) (io.ReadCloser, error) { + httpClient := *c.client + httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + } + + res, err := httpClient.Do(req.WithContext(ctx)) + if err != nil { + return nil, fmt.Errorf("initial signed download request: %w", err) + } + + switch { + case res.StatusCode >= 300 && res.StatusCode < 400: + location, err := res.Location() + res.Body.Close() + if err != nil { + return nil, fmt.Errorf("parsing signed download location: %w", err) + } + + blobReq, err := http.NewRequestWithContext(ctx, http.MethodGet, location.String(), nil) + if err != nil { + return nil, fmt.Errorf("creating signed download request: %w", err) + } + blobReq.Header.Set("User-Agent", c.UserAgent) + + res, err = cleanhttp.DefaultClient().Do(blobReq) + if err != nil { + return nil, fmt.Errorf("signed download request: %w", err) + } + if res.StatusCode >= 300 { + res.Body.Close() + return nil, fmt.Errorf("signed download returned %s", http.StatusText(res.StatusCode)) + } + case res.StatusCode >= 400: + defer res.Body.Close() + return nil, c.handleResponse(ctx, res, nil) + } + + return res.Body, nil +} From 0f8ce573f94f5209f6f2e838d5ddb9316295fe57 Mon Sep 17 00:00:00 2001 From: bwarminski Date: Wed, 29 Jul 2026 14:59:26 -0400 Subject: [PATCH 2/3] Add auth attempt export download command --- internal/cmd/auditlog/auditlogs.go | 4 +- internal/cmd/auditlog/auth_attempt_input.go | 192 ++++++++++ internal/cmd/auditlog/auth_attempts.go | 334 ++++++++++++++++++ internal/cmd/auditlog/auth_attempts_test.go | 293 +++++++++++++++ internal/cmd/branch/query_patterns.go | 3 +- internal/cmdutil/errors.go | 4 +- .../planetscale/auth_attempt_exports_test.go | 181 ---------- 7 files changed, 826 insertions(+), 185 deletions(-) create mode 100644 internal/cmd/auditlog/auth_attempt_input.go create mode 100644 internal/cmd/auditlog/auth_attempts.go create mode 100644 internal/cmd/auditlog/auth_attempts_test.go delete mode 100644 internal/planetscale/auth_attempt_exports_test.go diff --git a/internal/cmd/auditlog/auditlogs.go b/internal/cmd/auditlog/auditlogs.go index afeb9e9e..816a91e7 100644 --- a/internal/cmd/auditlog/auditlogs.go +++ b/internal/cmd/auditlog/auditlogs.go @@ -19,7 +19,8 @@ import ( func AuditLogCmd(ch *cmdutil.Helper) *cobra.Command { cmd := &cobra.Command{ Use: "audit-log ", - Short: "List audit logs", + Short: "List audit logs and download authentication-attempt exports", + Long: "List audit logs and download authentication-attempt export reports.", PersistentPreRunE: cmdutil.CheckAuthentication(ch.Config), } @@ -27,6 +28,7 @@ func AuditLogCmd(ch *cmdutil.Helper) *cobra.Command { cmd.MarkPersistentFlagRequired("org") // nolint:errcheck cmd.AddCommand(ListCmd(ch)) + cmd.AddCommand(AuthAttemptsCmd(ch)) return cmd } diff --git a/internal/cmd/auditlog/auth_attempt_input.go b/internal/cmd/auditlog/auth_attempt_input.go new file mode 100644 index 00000000..14c07b0c --- /dev/null +++ b/internal/cmd/auditlog/auth_attempt_input.go @@ -0,0 +1,192 @@ +package auditlog + +import ( + "fmt" + "slices" + "strconv" + "strings" + "time" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" + "github.com/spf13/cobra" +) + +var authAttemptExportFormats = []string{"jsonl", "csv", "parquet"} +var authAttemptExportOutcomes = []string{"allow", "deny"} +var authAttemptExportFailureReasons = []string{"bad_password", "unknown_user", "authorization_failed", "ip_not_allowed", "other"} +var authAttemptExportBackendRoutes = []string{"postgres", "pgbouncer", "unknown"} + +type downloadAuthAttemptsOptions struct { + since string + startAt string + endAt string + sourceIPs []string + branches []string + outcomes []string + usernames []string + startupDatabases []string + failureReasons []string + backendRoutes []string + exportFormat string + output string +} + +func authAttemptExportFormat(cmd *cobra.Command, requested string, output printer.Format) (string, error) { + if requested == "" { + if cmd.Flags().Changed("export-format") { + return "", fmt.Errorf("--export-format cannot be empty") + } + switch output { + case printer.JSON: + return "jsonl", nil + case printer.Human, printer.CSV: + return "csv", nil + default: + return "csv", nil + } + } + if !slices.Contains(authAttemptExportFormats, requested) { + return "", fmt.Errorf("invalid --export-format %q, must be one of: %s", requested, strings.Join(authAttemptExportFormats, ", ")) + } + return requested, nil +} + +func resolveAuthAttemptExportWindow(flags downloadAuthAttemptsOptions, now time.Time, location *time.Location) (time.Time, time.Time, error) { + switch { + case flags.since != "" && (flags.startAt != "" || flags.endAt != ""): + return time.Time{}, time.Time{}, fmt.Errorf("--since is mutually exclusive with --start-at and --end-at") + case flags.since == "" && flags.startAt == "" && flags.endAt == "": + return time.Time{}, time.Time{}, fmt.Errorf("a time window is required: use --since or --start-at with optional --end-at") + case flags.since == "" && flags.startAt == "": + return time.Time{}, time.Time{}, fmt.Errorf("--start-at is required when --since is not set") + } + + if flags.since != "" { + duration, err := parseAuthAttemptExportSince(flags.since) + if err != nil { + return time.Time{}, time.Time{}, err + } + return now.Add(-duration).UTC(), now.UTC(), nil + } + + startAt, err := parseAuthAttemptExportTime("start-at", flags.startAt, now, location) + if err != nil { + return time.Time{}, time.Time{}, err + } + endAt := now + if flags.endAt != "" { + endAt, err = parseAuthAttemptExportTime("end-at", flags.endAt, now, location) + if err != nil { + return time.Time{}, time.Time{}, err + } + } + if !startAt.Before(endAt) { + return time.Time{}, time.Time{}, fmt.Errorf("--start-at must be before --end-at") + } + return startAt.UTC(), endAt.UTC(), nil +} + +func parseAuthAttemptExportSince(value string) (time.Duration, error) { + if value == "" { + return 0, fmt.Errorf("--since cannot be empty; use a positive duration such as 24h, 7d, or 2w") + } + + const week = 7 * 24 * time.Hour + const maxDuration = time.Duration(1<<63 - 1) + if strings.HasSuffix(value, "w") { + weeks, err := strconv.ParseInt(value[:len(value)-1], 10, 64) + if err != nil || weeks <= 0 || weeks > int64(maxDuration/week) { + return 0, fmt.Errorf("invalid --since %q: use a positive duration such as 24h, 7d, or 2w", value) + } + return time.Duration(weeks) * week, nil + } + + duration, err := cmdutil.ParseDuration(value) + if err != nil { + return 0, fmt.Errorf("invalid --since %q: use a positive duration such as 24h, 7d, or 2w", value) + } + if strings.HasSuffix(value, "d") { + days, parseErr := strconv.ParseInt(value[:len(value)-1], 10, 64) + if parseErr != nil || days <= 0 || days > int64(maxDuration/(24*time.Hour)) { + return 0, fmt.Errorf("invalid --since %q: use a positive duration such as 24h, 7d, or 2w", value) + } + } + if duration <= 0 { + return 0, fmt.Errorf("invalid --since %q: use a positive duration such as 24h, 7d, or 2w", value) + } + return duration, nil +} + +func parseAuthAttemptExportTime(flag, value string, now time.Time, location *time.Location) (time.Time, error) { + switch strings.ToLower(value) { + case "now": + return now.UTC(), nil + case "today", "yesterday": + localNow := now.In(location) + midnight := time.Date(localNow.Year(), localNow.Month(), localNow.Day(), 0, 0, 0, 0, location) + if strings.EqualFold(value, "yesterday") { + midnight = midnight.AddDate(0, 0, -1) + } + return midnight.UTC(), nil + } + + if parsed, err := time.Parse(time.RFC3339, value); err == nil { + return parsed.UTC(), nil + } + for _, layout := range []string{"2006-01-02", "2006-01-02T15:04", "2006-01-02 15:04", "2006-01-02T15:04:05", "2006-01-02 15:04:05"} { + if parsed, err := time.ParseInLocation(layout, value, location); err == nil { + return parsed.UTC(), nil + } + } + return time.Time{}, fmt.Errorf("invalid --%s %q: use now, today, yesterday, YYYY-MM-DD, local ISO, or RFC3339", flag, value) +} + +func authAttemptExportFilters(sourceIPs, branches, outcomes, usernames, startupDatabases, failureReasons, backendRoutes []string, present map[string]bool) (ps.AuthAttemptExportFilters, error) { + if err := validateAuthAttemptFilter("source-ip", sourceIPs, nil, present["source-ip"]); err != nil { + return ps.AuthAttemptExportFilters{}, err + } + if err := validateAuthAttemptFilter("branch", branches, nil, present["branch"]); err != nil { + return ps.AuthAttemptExportFilters{}, err + } + if err := validateAuthAttemptFilter("outcome", outcomes, authAttemptExportOutcomes, present["outcome"]); err != nil { + return ps.AuthAttemptExportFilters{}, err + } + if err := validateAuthAttemptFilter("username", usernames, nil, present["username"]); err != nil { + return ps.AuthAttemptExportFilters{}, err + } + if err := validateAuthAttemptFilter("failure-reason", failureReasons, authAttemptExportFailureReasons, present["failure-reason"]); err != nil { + return ps.AuthAttemptExportFilters{}, err + } + if err := validateAuthAttemptFilter("backend-route", backendRoutes, authAttemptExportBackendRoutes, present["backend-route"]); err != nil { + return ps.AuthAttemptExportFilters{}, err + } + + return ps.AuthAttemptExportFilters{ + SourceIPs: sourceIPs, + Branches: branches, + Outcomes: outcomes, + Usernames: usernames, + StartupDatabases: startupDatabases, + FailureReasons: failureReasons, + BackendRoutes: backendRoutes, + }, nil +} + +func validateAuthAttemptFilter(flag string, values, allowed []string, changed bool) error { + if changed && len(values) == 0 { + return fmt.Errorf("--%s cannot be empty", flag) + } + for _, value := range values { + if value == "" { + return fmt.Errorf("--%s cannot be empty", flag) + } + if len(allowed) > 0 { + if !slices.Contains(allowed, value) { + return fmt.Errorf("invalid --%s %q, must be one of: %s", flag, value, strings.Join(allowed, ", ")) + } + } + } + return nil +} diff --git a/internal/cmd/auditlog/auth_attempts.go b/internal/cmd/auditlog/auth_attempts.go new file mode 100644 index 00000000..96384a35 --- /dev/null +++ b/internal/cmd/auditlog/auth_attempts.go @@ -0,0 +1,334 @@ +package auditlog + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "time" + + "github.com/planetscale/cli/internal/cmdutil" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" + "github.com/spf13/cobra" +) + +var authAttemptExportPollInterval = 2 * time.Second + +// AuthAttemptsCmd wraps commands for authentication-attempt exports. +func AuthAttemptsCmd(ch *cmdutil.Helper) *cobra.Command { + cmd := &cobra.Command{ + Use: "auth-attempts ", + Short: "Download authentication-attempt export reports", + } + cmd.AddCommand(DownloadAuthAttemptsCmd(ch)) + return cmd +} + +// AuthAttemptExportDownload is the result of downloading an authentication-attempt export. +type AuthAttemptExportDownload struct { + ID string `header:"id" json:"id"` + State string `header:"state" json:"state"` + Format string `header:"format" json:"format"` + StartAt string `header:"start_at" json:"start_at"` + EndAt string `header:"end_at" json:"end_at"` + File string `header:"file" json:"file"` +} + +func (a *AuthAttemptExportDownload) MarshalCSVValue() interface{} { + return []*AuthAttemptExportDownload{a} +} + +// DownloadAuthAttemptsCmd generates an authentication-attempt export, waits for it to finish, and downloads the ZIP file. +func DownloadAuthAttemptsCmd(ch *cmdutil.Helper) *cobra.Command { + flags := downloadAuthAttemptsOptions{} + + cmd := &cobra.Command{ + Use: "download", + Short: "Download an authentication-attempt export", + Long: "Generate an authentication-attempt export for a UTC time window, wait for it to be ready, and download its ZIP artifact. Use --since with a positive duration (24h, 7d, or 2w), or use --start-at and optionally --end-at. Named values and zone-less timestamps use the local timezone; requests are sent in UTC.", + Example: ` # Export the last 24 hours. + pscale audit-log auth-attempts download --since 24h + + # Export from local midnight yesterday through now. + pscale audit-log auth-attempts download --start-at yesterday + + # Export a date-only window. + pscale audit-log auth-attempts download --start-at 2026-07-28 --end-at 2026-07-29 + + # Zone-less local ISO timestamps can use T or a space. + pscale audit-log auth-attempts download --start-at '2026-07-29 16:00' --end-at '2026-07-29 17:00' + + # RFC3339 timestamps may include any explicit offset. + pscale audit-log auth-attempts download \ + --start-at 2026-07-29T00:00:00Z --end-at 2026-07-29T01:00:00Z + + # Export denied attempts from two addresses during a UTC window. + pscale audit-log auth-attempts download \ + --start-at 2026-07-29T00:00:00Z --end-at 2026-07-29T01:00:00Z \ + --source-ip 203.0.113.0/24 --source-ip 2001:db8::/32`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, args []string) error { + return runDownloadAuthAttempts(cmd, ch, flags) + }, + } + + cmd.Flags().StringVar(&flags.since, "since", "", "Export window ending now: positive duration such as 24h, 7d, or 2w") + cmd.Flags().StringVar(&flags.startAt, "start-at", "", "Inclusive start: now, today, yesterday, YYYY-MM-DD, local ISO, or RFC3339") + cmd.Flags().StringVar(&flags.endAt, "end-at", "", "Exclusive end: now, today, yesterday, YYYY-MM-DD, local ISO, or RFC3339; defaults to now") + cmd.Flags().StringSliceVar(&flags.sourceIPs, "source-ip", nil, "Source IP address or CIDR range (repeat or comma-separate)") + cmd.Flags().StringSliceVar(&flags.branches, "branch", nil, "Branch public ID or database/branch name (repeat or comma-separate)") + cmd.Flags().StringSliceVar(&flags.outcomes, "outcome", nil, fmt.Sprintf("Authentication outcome: %s", strings.Join(authAttemptExportOutcomes, " or "))) + cmd.Flags().StringArrayVar(&flags.usernames, "username", nil, "Authentication username (repeatable; commas are part of the name, not separators)") + cmd.Flags().StringArrayVar(&flags.startupDatabases, "startup-database", nil, "Startup database name (repeatable)") + cmd.Flags().StringSliceVar(&flags.failureReasons, "failure-reason", nil, fmt.Sprintf("Failure reason: %s", strings.Join(authAttemptExportFailureReasons, ", "))) + cmd.Flags().StringSliceVar(&flags.backendRoutes, "backend-route", nil, fmt.Sprintf("Backend route: %s", strings.Join(authAttemptExportBackendRoutes, ", "))) + cmd.Flags().StringVar(&flags.exportFormat, "export-format", "", fmt.Sprintf("Artifact format: %s; defaults from --format (human/csv -> csv, json -> JSONL; parquet is explicit)", strings.Join(authAttemptExportFormats, ", "))) + cmd.Flags().StringVar(&flags.output, "output", "", "Output file name, or - to write raw ZIP bytes to stdout") + + for _, enum := range []struct { + name string + values []string + }{ + {name: "outcome", values: authAttemptExportOutcomes}, + {name: "failure-reason", values: authAttemptExportFailureReasons}, + {name: "backend-route", values: authAttemptExportBackendRoutes}, + {name: "export-format", values: authAttemptExportFormats}, + } { + _ = cmd.RegisterFlagCompletionFunc(enum.name, cobra.FixedCompletions(enum.values, cobra.ShellCompDirectiveNoFileComp)) + } + cmd.MarkFlagsMutuallyExclusive("since", "start-at") + cmd.MarkFlagsMutuallyExclusive("since", "end-at") + + return cmd +} + +func runDownloadAuthAttempts(cmd *cobra.Command, ch *cmdutil.Helper, flags downloadAuthAttemptsOptions) error { + if cmd.Flags().Changed("since") && flags.since == "" { + return fmt.Errorf("--since cannot be empty; use a positive duration such as 24h, 7d, or 2w") + } + if cmd.Flags().Changed("start-at") && flags.startAt == "" { + return fmt.Errorf("--start-at cannot be empty; use now, today, yesterday, YYYY-MM-DD, local ISO, or RFC3339") + } + if cmd.Flags().Changed("end-at") && flags.endAt == "" { + return fmt.Errorf("--end-at cannot be empty; use now, today, yesterday, YYYY-MM-DD, local ISO, or RFC3339") + } + + startAt, endAt, err := resolveAuthAttemptExportWindow(flags, time.Now(), time.Local) + if err != nil { + return err + } + + filters, err := authAttemptExportFilters(flags.sourceIPs, flags.branches, flags.outcomes, + flags.usernames, flags.startupDatabases, flags.failureReasons, flags.backendRoutes, map[string]bool{ + "source-ip": cmd.Flags().Changed("source-ip"), + "branch": cmd.Flags().Changed("branch"), + "outcome": cmd.Flags().Changed("outcome"), + "username": cmd.Flags().Changed("username"), + "failure-reason": cmd.Flags().Changed("failure-reason"), + "backend-route": cmd.Flags().Changed("backend-route"), + }) + if err != nil { + return err + } + format, err := authAttemptExportFormat(cmd, flags.exportFormat, ch.Printer.Format()) + if err != nil { + return err + } + + client, err := ch.Client() + if err != nil { + return err + } + + toStdout := flags.output == "-" + endProgress := func() {} + if !toStdout { + endProgress = ch.Printer.PrintProgress(fmt.Sprintf("Generating authentication-attempt export for %s...", + printer.BoldBlue(ch.Config.Organization))) + } + defer endProgress() + + report, exportID, err := createAndPollAuthAttemptExport(cmd.Context(), client, ch.Config.Organization, + startAt, endAt, format, filters) + if err != nil { + return err + } + if report.State == "failed" { + return authAttemptExportFailureError(report, exportID) + } + if report.State != "ready" { + return fmt.Errorf("auth attempt export %s reached unexpected state %q", printer.BoldBlue(exportID), report.State) + } + + path := authAttemptExportOutputPath(flags.output, ch.Config.Organization, format) + if err := downloadAndPrintAuthAttemptExport(cmd, ch, client, exportID, format, path, + startAt, endAt, toStdout, report.State, endProgress); err != nil { + return err + } + return nil +} + +func createAndPollAuthAttemptExport(ctx context.Context, client *ps.Client, organization string, startAt, endAt time.Time, format string, filters ps.AuthAttemptExportFilters) (*ps.AuthAttemptExport, string, error) { + report, err := client.AuthAttemptExports.CreateExport(ctx, &ps.CreateAuthAttemptExportRequest{ + Organization: organization, + StartAt: startAt, + EndAt: endAt, + Format: format, + Filters: filters, + }) + if err != nil { + if cmdutil.ErrCode(err) == ps.ErrNotFound { + return nil, "", authAttemptExportNotFoundError(organization, "") + } + return nil, "", cmdutil.HandleError(err) + } + exportID := report.PublicID + for report.State == "pending" || report.State == "running" { + interval := report.RetryAfter + if interval <= 0 { + interval = authAttemptExportPollInterval + } + if err := waitForAuthAttemptExport(ctx, interval); err != nil { + return nil, exportID, authAttemptExportInterrupted(organization, exportID, err) + } + report, err = client.AuthAttemptExports.GetExport(ctx, &ps.GetAuthAttemptExportRequest{ + Organization: organization, Export: exportID, + }) + if err != nil { + if ctx.Err() != nil { + return nil, exportID, authAttemptExportInterrupted(organization, exportID, ctx.Err()) + } + if cmdutil.ErrCode(err) == ps.ErrNotFound { + return nil, exportID, authAttemptExportNotFoundError(organization, exportID) + } + return nil, exportID, cmdutil.HandleError(err) + } + } + return report, exportID, nil +} + +func downloadAndPrintAuthAttemptExport(cmd *cobra.Command, ch *cmdutil.Helper, client *ps.Client, exportID, format, path string, startAt, endAt time.Time, toStdout bool, state string, endProgress func()) error { + body, err := client.AuthAttemptExports.DownloadExport(cmd.Context(), &ps.DownloadAuthAttemptExportRequest{ + Organization: ch.Config.Organization, + Export: exportID, + }) + if err != nil { + if cmd.Context().Err() != nil { + return authAttemptExportInterrupted(ch.Config.Organization, exportID, cmd.Context().Err()) + } + if cmdutil.ErrCode(err) == ps.ErrNotFound { + return authAttemptExportNotFoundError(ch.Config.Organization, exportID) + } + return cmdutil.HandleError(err) + } + + if toStdout { + defer body.Close() + if _, err := io.Copy(cmd.OutOrStdout(), body); err != nil { + if cmd.Context().Err() != nil { + return authAttemptExportInterrupted(ch.Config.Organization, exportID, cmd.Context().Err()) + } + return fmt.Errorf("writing auth attempt export to stdout: %w", err) + } + return nil + } + if err := publishAuthAttemptExport(cmd.Context(), body, path); err != nil { + if cmd.Context().Err() != nil { + return authAttemptExportInterrupted(ch.Config.Organization, exportID, cmd.Context().Err()) + } + return fmt.Errorf("writing auth attempt export to %s: %w", path, err) + } + + endProgress() + if ch.Printer.Format() == printer.Human { + ch.Printer.Printf("Successfully downloaded auth attempts (%s) to %s\n", format, printer.BoldBlue(path)) + return nil + } + return ch.Printer.PrintResource(&AuthAttemptExportDownload{ + ID: exportID, State: state, Format: format, + StartAt: startAt.Format(time.RFC3339), EndAt: endAt.Format(time.RFC3339), File: path, + }) +} + +func waitForAuthAttemptExport(ctx context.Context, interval time.Duration) error { + timer := time.NewTimer(interval) + defer func() { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + }() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func authAttemptExportOutputPath(output, organization, format string) string { + filename := fmt.Sprintf("auth-attempts-%s-%s-%s.zip", organization, format, time.Now().UTC().Format("20060102T150405Z")) + if output == "" { + return filename + } + return output +} + +func publishAuthAttemptExport(ctx context.Context, body io.ReadCloser, path string) error { + tmp, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".tmp-*") + if err != nil { + body.Close() + return fmt.Errorf("creating temporary file: %w", err) + } + tmpPath := tmp.Name() + defer os.Remove(tmpPath) + + _, copyErr := io.Copy(tmp, body) + bodyCloseErr := body.Close() + closeErr := tmp.Close() + if copyErr == nil { + copyErr = bodyCloseErr + } + if copyErr == nil { + copyErr = closeErr + } + if copyErr != nil { + return copyErr + } + if err := ctx.Err(); err != nil { + return err + } + if err := os.Rename(tmpPath, path); err != nil { + return fmt.Errorf("renaming temporary file: %w", err) + } + return nil +} + +func authAttemptExportInterrupted(organization, exportID string, err error) error { + return fmt.Errorf("auth attempt export %s interrupted: %w; recover with pscale api organizations/%s/auth-attempt-exports/%s", + printer.BoldBlue(exportID), err, organization, exportID) +} + +func authAttemptExportFailureError(report *ps.AuthAttemptExport, exportID string) error { + message := fmt.Sprintf("auth attempt export %s failed (%s)", printer.BoldBlue(exportID), report.FailureReason) + if report.FailureDetail != "" { + message += ": " + report.FailureDetail + } + if report.RecoveryHint != "" { + message += "\n" + report.RecoveryHint + } + return errors.New(message) +} + +func authAttemptExportNotFoundError(organization, exportID string) error { + if exportID == "" { + return fmt.Errorf("auth-attempt exports are unavailable for organization %s", printer.BoldBlue(organization)) + } + return fmt.Errorf("auth-attempt export %s was not found and may have expired", printer.BoldBlue(exportID)) +} diff --git a/internal/cmd/auditlog/auth_attempts_test.go b/internal/cmd/auditlog/auth_attempts_test.go new file mode 100644 index 00000000..2541b81d --- /dev/null +++ b/internal/cmd/auditlog/auth_attempts_test.go @@ -0,0 +1,293 @@ +package auditlog + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + qt "github.com/frankban/quicktest" + + "github.com/planetscale/cli/internal/cmdutil" + "github.com/planetscale/cli/internal/config" + ps "github.com/planetscale/cli/internal/planetscale" + "github.com/planetscale/cli/internal/printer" +) + +func authAttemptTestHelper(org, baseURL string, format printer.Format, buf *bytes.Buffer) *cmdutil.Helper { + p := printer.NewPrinter(&format) + p.SetResourceOutput(buf) + p.SetHumanOutput(io.Discard) + + return &cmdutil.Helper{ + Printer: p, + Config: &config.Config{AccessToken: "token", Organization: org, BaseURL: baseURL}, + Client: func() (*ps.Client, error) { + return ps.NewClient(ps.WithBaseURL(baseURL), ps.WithAccessToken("token")) + }, + } +} + +func TestAuthAttemptsDownloadCreatesPollsAndWritesFile(t *testing.T) { + c := qt.New(t) + + previousInterval := authAttemptExportPollInterval + authAttemptExportPollInterval = 5 * time.Second + t.Cleanup(func() { authAttemptExportPollInterval = previousInterval }) + + var createRequests, statusRequests, downloadRequests int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/v1/organizations/my-org/auth-attempt-exports": + createRequests++ + var body struct { + StartAt string `json:"start_at"` + EndAt string `json:"end_at"` + Format string `json:"format"` + Filters struct { + SourceIPs []string `json:"source_ips"` + Branches []string `json:"branches"` + Outcomes []string `json:"outcomes"` + Usernames []string `json:"usernames"` + StartupDatabases []string `json:"startup_databases"` + FailureReasons []string `json:"failure_reasons"` + BackendRoutes []string `json:"backend_routes"` + } `json:"filters"` + } + c.Assert(json.NewDecoder(r.Body).Decode(&body), qt.IsNil) + c.Assert(body.StartAt, qt.Equals, "2026-07-29T00:00:00Z") + c.Assert(body.EndAt, qt.Equals, "2026-07-29T01:00:00Z") + c.Assert(body.Format, qt.Equals, "jsonl") + c.Assert(body.Filters.SourceIPs, qt.DeepEquals, []string{"203.0.113.0/24", "2001:db8::/112"}) + c.Assert(body.Filters.Branches, qt.DeepEquals, []string{"db/production"}) + c.Assert(body.Filters.Outcomes, qt.DeepEquals, []string{"deny"}) + c.Assert(body.Filters.Usernames, qt.DeepEquals, []string{"incident-user"}) + c.Assert(body.Filters.StartupDatabases, qt.DeepEquals, []string{""}) + c.Assert(body.Filters.FailureReasons, qt.DeepEquals, []string{"bad_password"}) + c.Assert(body.Filters.BackendRoutes, qt.DeepEquals, []string{"postgres"}) + w.Header().Set("Retry-After", "1") + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"export1","state":"pending","format":"jsonl"}`) + case r.Method == http.MethodGet && r.URL.Path == "/v1/organizations/my-org/auth-attempt-exports/export1": + statusRequests++ + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"export1","state":"ready","start_at":"2026-07-29T00:00:00Z","end_at":"2026-07-29T01:00:00Z","format":"jsonl"}`) + case r.Method == http.MethodGet && r.URL.Path == "/v1/organizations/my-org/auth-attempt-exports/export1/download": + downloadRequests++ + http.Redirect(w, r, "/blob", http.StatusFound) + case r.Method == http.MethodGet && r.URL.Path == "/blob": + _, _ = io.WriteString(w, "PK\x03\x04zip") + default: + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + + output := filepath.Join(t.TempDir(), "auth-attempts.zip") + var printed bytes.Buffer + cmd := AuthAttemptsCmd(authAttemptTestHelper("my-org", server.URL, printer.JSON, &printed)) + cmd.SetErr(io.Discard) + cmd.SetArgs([]string{ + "download", + "--start-at", "2026-07-29T00:00:00Z", + "--end-at", "2026-07-29T01:00:00Z", + "--source-ip", "203.0.113.0/24,2001:db8::/112", + "--branch", "db/production", + "--outcome", "deny", + "--username", "incident-user", + "--startup-database", "", + "--failure-reason", "bad_password", + "--backend-route", "postgres", + "--output", output, + }) + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + t.Cleanup(cancel) + c.Assert(cmd.ExecuteContext(ctx), qt.IsNil) + c.Assert(createRequests, qt.Equals, 1) + c.Assert(statusRequests, qt.Equals, 1) + c.Assert(downloadRequests, qt.Equals, 1) + + downloaded, err := os.ReadFile(output) + c.Assert(err, qt.IsNil) + c.Assert(downloaded, qt.DeepEquals, []byte("PK\x03\x04zip")) + c.Assert(printed.String(), qt.JSONEquals, &AuthAttemptExportDownload{ + ID: "export1", State: "ready", Format: "jsonl", + StartAt: "2026-07-29T00:00:00Z", EndAt: "2026-07-29T01:00:00Z", File: output, + }) +} + +func TestAuthAttemptsDownloadHonorsContextCancellation(t *testing.T) { + for _, test := range []struct { + name string + ctx func() (context.Context, context.CancelFunc) + want error + }{ + {name: "canceled", ctx: func() (context.Context, context.CancelFunc) { + return context.WithCancel(context.Background()) + }, want: context.Canceled}, + {name: "deadline exceeded", ctx: func() (context.Context, context.CancelFunc) { + return context.WithTimeout(context.Background(), 20*time.Millisecond) + }, want: context.DeadlineExceeded}, + } { + t.Run(test.name, func(t *testing.T) { + c := qt.New(t) + previousInterval := authAttemptExportPollInterval + authAttemptExportPollInterval = time.Hour + t.Cleanup(func() { authAttemptExportPollInterval = previousInterval }) + + created := make(chan struct{}) + var downloads int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost: + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"export1","state":"pending","format":"jsonl"}`) + close(created) + case r.Method == http.MethodGet && r.URL.Path == "/v1/organizations/my-org/auth-attempt-exports/export1/download": + downloads++ + http.Error(w, "unexpected download", http.StatusInternalServerError) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + + ctx, cancel := test.ctx() + t.Cleanup(cancel) + var printed bytes.Buffer + cmd := AuthAttemptsCmd(authAttemptTestHelper("my-org", server.URL, printer.JSON, &printed)) + cmd.SetErr(io.Discard) + cmd.SetArgs([]string{ + "download", + "--start-at", "2026-07-29T00:00:00Z", + "--end-at", "2026-07-29T01:00:00Z", + "--output", filepath.Join(t.TempDir(), "auth-attempts.zip"), + }) + + result := make(chan error, 1) + go func() { result <- cmd.ExecuteContext(ctx) }() + select { + case <-created: + case <-time.After(time.Second): + t.Fatal("export was not created") + } + if test.name == "canceled" { + time.Sleep(50 * time.Millisecond) + cancel() + } + err := <-result + c.Assert(errors.Is(err, test.want), qt.IsTrue) + c.Assert(err, qt.ErrorMatches, `.*export1.*`) + c.Assert(err, qt.ErrorMatches, `.*pscale api organizations/my-org/auth-attempt-exports/export1.*`) + c.Assert(downloads, qt.Equals, 0) + }) + } +} + +func TestAuthAttemptsDownloadFilterPresence(t *testing.T) { + tests := []struct { + name string + args []string + wantAbsent []string + want map[string][]string + }{ + {name: "absent", wantAbsent: []string{ + "source_ips", "branches", "outcomes", "usernames", "startup_databases", "failure_reasons", "backend_routes", + }}, + {name: "empty startup database", args: []string{"--startup-database", ""}, want: map[string][]string{"startup_databases": {""}}}, + {name: "startup database order", args: []string{"--startup-database", "", "--startup-database", "postgres"}, want: map[string][]string{"startup_databases": {"", "postgres"}}}, + {name: "username commas stay together", args: []string{"--username", "a,b"}, want: map[string][]string{"usernames": {"a,b"}}}, + {name: "source and enum comma values", args: []string{"--source-ip", "203.0.113.0/24,2001:db8::/112", "--outcome", "allow,deny", "--failure-reason", "bad_password,other", "--backend-route", "postgres,unknown"}, want: map[string][]string{ + "source_ips": {"203.0.113.0/24", "2001:db8::/112"}, "outcomes": {"allow", "deny"}, "failure_reasons": {"bad_password", "other"}, "backend_routes": {"postgres", "unknown"}, + }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + c := qt.New(t) + var posted map[string]json.RawMessage + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + var body struct { + Filters map[string]json.RawMessage `json:"filters"` + } + c.Assert(json.NewDecoder(r.Body).Decode(&body), qt.IsNil) + posted = body.Filters + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"export1","state":"ready","format":"jsonl","start_at":"2026-07-29T00:00:00Z","end_at":"2026-07-29T01:00:00Z"}`) + return + } + if r.URL.Path == "/v1/organizations/my-org/auth-attempt-exports/export1/download" { + http.Redirect(w, r, "/blob", http.StatusFound) + return + } + if r.URL.Path == "/blob" { + _, _ = io.WriteString(w, "PK\x03\x04zip") + return + } + http.NotFound(w, r) + })) + t.Cleanup(server.Close) + + args := append([]string{"download", "--start-at", "2026-07-29T00:00:00Z", "--end-at", "2026-07-29T01:00:00Z", "--output", filepath.Join(t.TempDir(), "auth-attempts.zip")}, test.args...) + var printed bytes.Buffer + cmd := AuthAttemptsCmd(authAttemptTestHelper("my-org", server.URL, printer.JSON, &printed)) + cmd.SetArgs(args) + c.Assert(cmd.Execute(), qt.IsNil) + + for key, values := range test.want { + var got []string + c.Assert(json.Unmarshal(posted[key], &got), qt.IsNil) + c.Assert(got, qt.DeepEquals, values) + } + for _, key := range test.wantAbsent { + _, ok := posted[key] + c.Assert(ok, qt.IsFalse, qt.Commentf("filter %s was sent", key)) + } + }) + } +} + +func TestAuthAttemptsDownloadInterruptedCopyDoesNotPublishPartialFile(t *testing.T) { + c := qt.New(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost: + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"export1","state":"ready","format":"jsonl"}`) + case strings.HasSuffix(r.URL.Path, "/download"): + http.Redirect(w, r, "/blob", http.StatusFound) + case r.URL.Path == "/blob": + _, _ = io.WriteString(w, "partial") + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + <-r.Context().Done() + default: + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + + output := filepath.Join(t.TempDir(), "auth-attempts.zip") + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + t.Cleanup(cancel) + cmd := AuthAttemptsCmd(authAttemptTestHelper("my-org", server.URL, printer.JSON, &bytes.Buffer{})) + cmd.SetErr(io.Discard) + cmd.SetArgs([]string{"download", "--start-at", "2026-07-29T00:00:00Z", "--end-at", "2026-07-29T01:00:00Z", "--output", output}) + c.Assert(cmd.ExecuteContext(ctx), qt.IsNotNil) + _, err := os.Stat(output) + c.Assert(err, qt.ErrorIs, os.ErrNotExist) + entries, err := os.ReadDir(filepath.Dir(output)) + c.Assert(err, qt.IsNil) + c.Assert(len(entries), qt.Equals, 0) +} diff --git a/internal/cmd/branch/query_patterns.go b/internal/cmd/branch/query_patterns.go index ba53c34c..52e6e79e 100644 --- a/internal/cmd/branch/query_patterns.go +++ b/internal/cmd/branch/query_patterns.go @@ -159,7 +159,8 @@ func DownloadQueryPatternsCmd(ch *cmdutil.Helper) *cobra.Command { // queryPatternsError maps a not-found API error to a message explaining both // causes: a missing branch, or query insights being disabled for the database. func queryPatternsError(ch *cmdutil.Helper, err error, database, branch string) error { - switch cmdutil.ErrCode(err) { + code := cmdutil.ErrCode(err) + switch code { case ps.ErrNotFound: 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)) diff --git a/internal/cmdutil/errors.go b/internal/cmdutil/errors.go index 2f3a800d..a7e06ab5 100644 --- a/internal/cmdutil/errors.go +++ b/internal/cmdutil/errors.go @@ -39,8 +39,8 @@ func ErrCode(err error) planetscale.ErrorCode { return "" } - perr, ok := err.(*planetscale.Error) - if !ok { + var perr *planetscale.Error + if !errors.As(err, &perr) { return "" } diff --git a/internal/planetscale/auth_attempt_exports_test.go b/internal/planetscale/auth_attempt_exports_test.go deleted file mode 100644 index 0cf10ba0..00000000 --- a/internal/planetscale/auth_attempt_exports_test.go +++ /dev/null @@ -1,181 +0,0 @@ -package planetscale - -import ( - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "testing" - "time" - - qt "github.com/frankban/quicktest" -) - -func TestAuthAttemptExports_CreateExport(t *testing.T) { - c := qt.New(t) - - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - c.Assert(r.Method, qt.Equals, http.MethodPost) - c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/auth-attempt-exports") - - var got map[string]any - c.Assert(json.NewDecoder(r.Body).Decode(&got), qt.IsNil) - want := map[string]any{ - "start_at": "2026-07-29T00:00:00Z", - "end_at": "2026-07-29T01:00:00Z", - "format": "parquet", - "filters": map[string]any{ - "source_ips": []any{"203.0.113.0/24", "2001:db8::/112"}, - "branches": []any{"db/production"}, - "outcomes": []any{"deny"}, - "usernames": []any{"incident-user"}, - "startup_databases": []any{""}, - "failure_reasons": []any{"bad_password"}, - "backend_routes": []any{"postgres"}, - }, - } - c.Assert(got, qt.DeepEquals, want) - - w.Header().Set("Retry-After", "5") - w.WriteHeader(http.StatusCreated) - _, err := w.Write([]byte(`{"id":"export1","state":"pending","format":"parquet"}`)) - c.Assert(err, qt.IsNil) - })) - t.Cleanup(ts.Close) - - client, err := NewClient(WithBaseURL(ts.URL)) - c.Assert(err, qt.IsNil) - - export, err := client.AuthAttemptExports.CreateExport(context.Background(), &CreateAuthAttemptExportRequest{ - Organization: "my-org", - StartAt: time.Date(2026, time.July, 29, 0, 0, 0, 0, time.UTC), - EndAt: time.Date(2026, time.July, 29, 1, 0, 0, 0, time.UTC), - Format: "parquet", - Filters: AuthAttemptExportFilters{ - SourceIPs: []string{"203.0.113.0/24", "2001:db8::/112"}, - Branches: []string{"db/production"}, - Outcomes: []string{"deny"}, - Usernames: []string{"incident-user"}, - StartupDatabases: []string{""}, - FailureReasons: []string{"bad_password"}, - BackendRoutes: []string{"postgres"}, - }, - }) - - c.Assert(err, qt.IsNil) - c.Assert(export.PublicID, qt.Equals, "export1") - c.Assert(export.State, qt.Equals, "pending") - c.Assert(export.Format, qt.Equals, "parquet") - c.Assert(export.RetryAfter, qt.Equals, 5*time.Second) -} - -func TestAuthAttemptExports_GetExport(t *testing.T) { - c := qt.New(t) - - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - c.Assert(r.Method, qt.Equals, http.MethodGet) - c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/auth-attempt-exports/export1") - - w.Header().Set("Retry-After", "2") - w.WriteHeader(http.StatusOK) - _, err := w.Write([]byte(`{ - "id":"export1", - "state":"failed", - "start_at":"2026-07-29T00:00:00Z", - "end_at":"2026-07-29T01:00:00Z", - "format":"parquet", - "filters":{ - "source_ips":["203.0.113.0/24"], - "branches":["db/production"], - "outcomes":["deny"], - "usernames":["incident-user"], - "startup_databases":[""], - "failure_reasons":["bad_password"], - "backend_routes":["postgres"] - }, - "resolved_branch_public_ids":["branch1"], - "created_at":"2026-07-29T01:01:00Z", - "started_at":"2026-07-29T01:02:00Z", - "generated_at":"2026-07-29T01:03:00Z", - "finished_at":"2026-07-29T01:04:00Z", - "expires_at":"2026-07-30T01:04:00Z", - "failure_reason":"generation_failed", - "failure_detail":"The export could not be generated.", - "recovery_hint":"Re-create the export; if the error repeats, contact support." - }`)) - c.Assert(err, qt.IsNil) - })) - t.Cleanup(ts.Close) - - client, err := NewClient(WithBaseURL(ts.URL)) - c.Assert(err, qt.IsNil) - - export, err := client.AuthAttemptExports.GetExport(context.Background(), &GetAuthAttemptExportRequest{ - Organization: "my-org", - Export: "export1", - }) - - startedAt := time.Date(2026, time.July, 29, 1, 2, 0, 0, time.UTC) - generatedAt := time.Date(2026, time.July, 29, 1, 3, 0, 0, time.UTC) - finishedAt := time.Date(2026, time.July, 29, 1, 4, 0, 0, time.UTC) - expiresAt := time.Date(2026, time.July, 30, 1, 4, 0, 0, time.UTC) - want := &AuthAttemptExport{ - PublicID: "export1", - State: "failed", - StartAt: time.Date(2026, time.July, 29, 0, 0, 0, 0, time.UTC), - EndAt: time.Date(2026, time.July, 29, 1, 0, 0, 0, time.UTC), - Format: "parquet", - Filters: AuthAttemptExportFilters{ - SourceIPs: []string{"203.0.113.0/24"}, - Branches: []string{"db/production"}, - Outcomes: []string{"deny"}, - Usernames: []string{"incident-user"}, - StartupDatabases: []string{""}, - FailureReasons: []string{"bad_password"}, - BackendRoutes: []string{"postgres"}, - }, - ResolvedBranchPublicIDs: []string{"branch1"}, - CreatedAt: time.Date(2026, time.July, 29, 1, 1, 0, 0, time.UTC), - StartedAt: &startedAt, - GeneratedAt: &generatedAt, - FinishedAt: &finishedAt, - ExpiresAt: &expiresAt, - FailureReason: "generation_failed", - FailureDetail: "The export could not be generated.", - RecoveryHint: "Re-create the export; if the error repeats, contact support.", - RetryAfter: 2 * time.Second, - } - - c.Assert(err, qt.IsNil) - c.Assert(export, qt.DeepEquals, want) -} - -func TestParseRetryAfter(t *testing.T) { - c := qt.New(t) - future := time.Now().Add(2 * time.Minute).UTC().Format(http.TimeFormat) - past := time.Now().Add(-time.Minute).UTC().Format(http.TimeFormat) - tests := []struct { - name string - header string - want time.Duration - }{ - {name: "positive delta seconds", header: "5", want: 5 * time.Second}, - {name: "future HTTP date", header: future}, - {name: "absent", header: "", want: 0}, - {name: "malformed", header: "soon", want: 0}, - {name: "fractional", header: "1.5", want: 0}, - {name: "zero", header: "0", want: 0}, - {name: "past HTTP date", header: past, want: 0}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - got := parseRetryAfter(test.header) - if test.name == "future HTTP date" { - c.Assert(got > 0, qt.IsTrue) - return - } - c.Assert(got, qt.Equals, test.want) - }) - } -} From 942dcfdc464d9af05c6dadbae8f6fb75da8d3a1a Mon Sep 17 00:00:00 2001 From: bwarminski Date: Wed, 29 Jul 2026 16:05:03 -0400 Subject: [PATCH 3/3] Handle wrapped API errors --- internal/cmdutil/errors.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/cmdutil/errors.go b/internal/cmdutil/errors.go index a7e06ab5..07fcb755 100644 --- a/internal/cmdutil/errors.go +++ b/internal/cmdutil/errors.go @@ -56,8 +56,8 @@ func HandleError(err error) error { return err } - perr, ok := err.(*planetscale.Error) - if !ok { + var perr *planetscale.Error + if !errors.As(err, &perr) { return err }