From bbbbeb7444b3addd9daf626b00f5a3133640f680 Mon Sep 17 00:00:00 2001 From: David Graham Date: Tue, 7 Jul 2026 13:01:49 -0600 Subject: [PATCH] Define query pattern report types Needed in https://github.com/planetscale/cli/pull/1281. --- planetscale/client.go | 2 + planetscale/query_patterns.go | 191 +++++++++++++++++++++++++++++ planetscale/query_patterns_test.go | 180 +++++++++++++++++++++++++++ 3 files changed, 373 insertions(+) create mode 100644 planetscale/query_patterns.go create mode 100644 planetscale/query_patterns_test.go diff --git a/planetscale/client.go b/planetscale/client.go index 3587de3..598aef5 100644 --- a/planetscale/client.go +++ b/planetscale/client.go @@ -66,6 +66,7 @@ type Client struct { PostgresBranches PostgresBranchesService PostgresRoles PostgresRolesService Processlist ProcesslistService + QueryPatterns QueryPatternsService Regions RegionsService SchemaRecommendations SchemaRecommendationService ServiceTokens ServiceTokenService @@ -306,6 +307,7 @@ func NewClient(opts ...ClientOption) (*Client, error) { c.Processlist = &processlistService{client: c} c.PostgresBranches = &postgresBranchesService{client: c} c.PostgresRoles = &postgresRolesService{client: c} + c.QueryPatterns = &queryPatternsService{client: c} c.Regions = ®ionsService{client: c} c.SchemaRecommendations = &schemaRecommendationService{client: c} c.ServiceTokens = &serviceTokenService{client: c} diff --git a/planetscale/query_patterns.go b/planetscale/query_patterns.go new file mode 100644 index 0000000..20b669d --- /dev/null +++ b/planetscale/query_patterns.go @@ -0,0 +1,191 @@ +package planetscale + +import ( + "bytes" + "compress/gzip" + "context" + "errors" + "fmt" + "io" + "net/http" + "path" + "time" + + "github.com/hashicorp/go-cleanhttp" +) + +// QueryPatternsReport represents a query patterns report for a branch. +type QueryPatternsReport struct { + PublicID string `json:"id"` + State string `json:"state"` + Actor *Actor `json:"actor"` + URL string `json:"url"` + DownloadURL string `json:"download_url"` + CreatedAt time.Time `json:"created_at"` + FinishedAt time.Time `json:"finished_at"` +} + +type CreateQueryPatternsReportRequest struct { + Organization string + Database string + Branch string +} + +type GetQueryPatternsReportRequest struct { + Organization string + Database string + Branch string + Report string +} + +type DownloadQueryPatternsReportRequest struct { + Organization string + Database string + Branch string + Report string +} + +// QueryPatternsService is an interface for communicating with the PlanetScale +// query patterns API endpoints. +type QueryPatternsService interface { + CreateReport(context.Context, *CreateQueryPatternsReportRequest) (*QueryPatternsReport, error) + GetReport(context.Context, *GetQueryPatternsReportRequest) (*QueryPatternsReport, error) + DownloadReport(context.Context, *DownloadQueryPatternsReportRequest) (io.ReadCloser, error) +} + +type queryPatternsService struct { + client *Client +} + +var _ QueryPatternsService = &queryPatternsService{} + +func NewQueryPatternsService(client *Client) *queryPatternsService { + return &queryPatternsService{ + client: client, + } +} + +// CreateReport starts generating a new query patterns report for a branch. +func (s *queryPatternsService) CreateReport(ctx context.Context, createReq *CreateQueryPatternsReportRequest) (*QueryPatternsReport, error) { + path := queryPatternsAPIPath(createReq.Organization, createReq.Database, createReq.Branch) + req, err := s.client.newRequest(http.MethodPost, path, nil) + if err != nil { + return nil, fmt.Errorf("error creating http request: %w", err) + } + + report := &QueryPatternsReport{} + if err := s.client.do(ctx, req, &report); err != nil { + return nil, err + } + + return report, nil +} + +// GetReport returns a single query patterns report for a branch. +func (s *queryPatternsService) GetReport(ctx context.Context, getReq *GetQueryPatternsReportRequest) (*QueryPatternsReport, error) { + path := queryPatternsReportAPIPath(getReq.Organization, getReq.Database, getReq.Branch, getReq.Report) + req, err := s.client.newRequest(http.MethodGet, path, nil) + if err != nil { + return nil, fmt.Errorf("error creating http request: %w", err) + } + + report := &QueryPatternsReport{} + if err := s.client.do(ctx, req, &report); err != nil { + return nil, err + } + + return report, nil +} + +// DownloadReport returns the content of a completed query patterns report. +// The caller must close the returned io.ReadCloser. +func (s *queryPatternsService) DownloadReport(ctx context.Context, downloadReq *DownloadQueryPatternsReportRequest) (io.ReadCloser, error) { + reqPath := path.Join(queryPatternsReportAPIPath(downloadReq.Organization, downloadReq.Database, downloadReq.Branch, downloadReq.Report), "download") + req, err := s.client.newRequest(http.MethodGet, reqPath, nil) + if err != nil { + 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)) + 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 decompressedReadCloser(res.Body) +} + +// decompressedReadCloser wraps body so gzip-compressed content is transparently +// decompressed, sniffing the gzip magic bytes so an uncompressed body passes +// through unchanged. +func decompressedReadCloser(body io.ReadCloser) (io.ReadCloser, error) { + head := make([]byte, 2) + n, err := io.ReadFull(body, head) + if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) { + body.Close() + return nil, err + } + + content := io.MultiReader(bytes.NewReader(head[:n]), body) + if n == 2 && head[0] == 0x1f && head[1] == 0x8b { + gz, err := gzip.NewReader(content) + if err != nil { + body.Close() + return nil, err + } + content = gz + } + + return &bodyReadCloser{Reader: content, body: body}, nil +} + +type bodyReadCloser struct { + io.Reader + body io.ReadCloser +} + +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") +} + +func queryPatternsReportAPIPath(org, db, branch, report string) string { + return path.Join(queryPatternsAPIPath(org, db, branch), report) +} diff --git a/planetscale/query_patterns_test.go b/planetscale/query_patterns_test.go new file mode 100644 index 0000000..9e715a7 --- /dev/null +++ b/planetscale/query_patterns_test.go @@ -0,0 +1,180 @@ +package planetscale + +import ( + "compress/gzip" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + qt "github.com/frankban/quicktest" +) + +const testQueryPatternsCSV = "normalized_sql,query_count\nselect ?,10\n" + +func TestQueryPatterns_CreateReport(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/databases/my-db/branches/my-branch/query-patterns") + w.WriteHeader(201) + out := `{"id":"report1","state":"pending","actor":{"id":"actor1","type":"User","display_name":"user@example.com"},"url":"https://api.planetscale.com/v1/organizations/my-org/databases/my-db/branches/my-branch/query-patterns/report1","created_at":"2021-01-14T10:19:23.000Z"}` + _, err := w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + report, err := client.QueryPatterns.CreateReport(context.Background(), &CreateQueryPatternsReportRequest{ + Organization: "my-org", + Database: "my-db", + Branch: "my-branch", + }) + + want := &QueryPatternsReport{ + PublicID: "report1", + State: "pending", + Actor: &Actor{ + ID: "actor1", + Type: "User", + Name: "user@example.com", + }, + URL: "https://api.planetscale.com/v1/organizations/my-org/databases/my-db/branches/my-branch/query-patterns/report1", + CreatedAt: time.Date(2021, time.January, 14, 10, 19, 23, 0, time.UTC), + } + + c.Assert(err, qt.IsNil) + c.Assert(report, qt.DeepEquals, want) +} + +func TestQueryPatterns_GetReport(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.Method, qt.Equals, http.MethodGet) + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/my-db/branches/my-branch/query-patterns/report1") + w.WriteHeader(200) + out := `{"id":"report1","state":"completed","download_url":"https://api.planetscale.com/v1/organizations/my-org/databases/my-db/branches/my-branch/query-patterns/report1/download","created_at":"2021-01-14T10:19:23.000Z","finished_at":"2021-01-14T10:20:23.000Z"}` + _, err := w.Write([]byte(out)) + c.Assert(err, qt.IsNil) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + report, err := client.QueryPatterns.GetReport(context.Background(), &GetQueryPatternsReportRequest{ + Organization: "my-org", + Database: "my-db", + Branch: "my-branch", + Report: "report1", + }) + + want := &QueryPatternsReport{ + PublicID: "report1", + State: "completed", + DownloadURL: "https://api.planetscale.com/v1/organizations/my-org/databases/my-db/branches/my-branch/query-patterns/report1/download", + CreatedAt: time.Date(2021, time.January, 14, 10, 19, 23, 0, time.UTC), + FinishedAt: time.Date(2021, time.January, 14, 10, 20, 23, 0, time.UTC), + } + + c.Assert(err, qt.IsNil) + c.Assert(report, qt.DeepEquals, want) +} + +func TestQueryPatterns_DownloadReport(t *testing.T) { + c := qt.New(t) + + // blob simulates presigned storage: it rejects requests that carry the + // API's Authorization header. + blob := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.Header.Get("Authorization"), qt.Equals, "") + w.Header().Set("Content-Type", "application/gzip") + gz := gzip.NewWriter(w) + _, err := gz.Write([]byte(testQueryPatternsCSV)) + c.Assert(err, qt.IsNil) + c.Assert(gz.Close(), qt.IsNil) + })) + t.Cleanup(blob.Close) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c.Assert(r.URL.Path, qt.Equals, "/v1/organizations/my-org/databases/my-db/branches/my-branch/query-patterns/report1/download") + c.Assert(r.Header.Get("Authorization"), qt.Equals, "tid:secret") + http.Redirect(w, r, blob.URL, http.StatusFound) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL), WithServiceToken("tid", "secret")) + c.Assert(err, qt.IsNil) + + body, err := client.QueryPatterns.DownloadReport(context.Background(), &DownloadQueryPatternsReportRequest{ + Organization: "my-org", + Database: "my-db", + Branch: "my-branch", + Report: "report1", + }) + c.Assert(err, qt.IsNil) + defer body.Close() + + data, err := io.ReadAll(body) + c.Assert(err, qt.IsNil) + c.Assert(string(data), qt.Equals, testQueryPatternsCSV) +} + +func TestQueryPatterns_DownloadReportUncompressed(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/csv") + _, err := w.Write([]byte(testQueryPatternsCSV)) + c.Assert(err, qt.IsNil) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + body, err := client.QueryPatterns.DownloadReport(context.Background(), &DownloadQueryPatternsReportRequest{ + Organization: "my-org", + Database: "my-db", + Branch: "my-branch", + Report: "report1", + }) + c.Assert(err, qt.IsNil) + defer body.Close() + + data, err := io.ReadAll(body) + c.Assert(err, qt.IsNil) + c.Assert(string(data), qt.Equals, testQueryPatternsCSV) +} + +func TestQueryPatterns_DownloadReportNotFound(t *testing.T) { + c := qt.New(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(404) + _, err := w.Write([]byte(`{"code":"not_found","message":"Not Found"}`)) + c.Assert(err, qt.IsNil) + })) + t.Cleanup(ts.Close) + + client, err := NewClient(WithBaseURL(ts.URL)) + c.Assert(err, qt.IsNil) + + _, err = client.QueryPatterns.DownloadReport(context.Background(), &DownloadQueryPatternsReportRequest{ + Organization: "my-org", + Database: "my-db", + Branch: "my-branch", + Report: "report1", + }) + c.Assert(err, qt.IsNotNil) + + var perr *Error + c.Assert(err, qt.ErrorAs, &perr) + c.Assert(perr.Code, qt.Equals, ErrNotFound) +}