Define query pattern report types

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 = &regionsService{client: c}
 	c.SchemaRecommendations = &schemaRecommendationService{client: c}
 	c.ServiceTokens = &serviceTokenService{client: c}

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)
+}

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)
+}