diff --git a/internal/corrosion/client.go b/internal/corrosion/client.go index 6e17ad6e..bee5d2d1 100644 --- a/internal/corrosion/client.go +++ b/internal/corrosion/client.go @@ -1,21 +1,17 @@ package corrosion import ( - "bytes" "context" "crypto/tls" - "encoding/json" "errors" "fmt" "github.com/cenkalti/backoff/v4" "golang.org/x/net/http2" - "io" "log/slog" "net" "net/http" "net/netip" "net/url" - "strconv" "time" ) @@ -124,252 +120,3 @@ func (rt *RetryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error boff := backoff.WithContext(rt.NewBackoff(), req.Context()) return backoff.RetryWithData(roundTrip, boff) } - -// SubscribeContext creates a subscription to receive updates for a desired SQL query. If skipRows is false, -// Subscription.Rows must be consumed before Subscription.Changes can be called. If skipRows is true, Subscription.Rows -// will be nil. -func (c *APIClient) SubscribeContext( - ctx context.Context, query string, args []any, skipRows bool, -) (*Subscription, error) { - statement := Statement{ - Query: query, - Params: args, - } - body, err := json.Marshal(statement) - if err != nil { - return nil, fmt.Errorf("marshal query: %w", err) - } - - subURL := c.baseURL.JoinPath("/v1/subscriptions") - if skipRows { - q := subURL.Query() - q.Set("skip_rows", "true") - subURL.RawQuery = q.Encode() - } - - req, err := http.NewRequestWithContext(ctx, "POST", subURL.String(), bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Accept", "application/json") - - resp, err := c.client.Do(req) - if err != nil { - return nil, fmt.Errorf("send request: %w", err) - } - - if resp.StatusCode != http.StatusOK { - respBody, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return nil, fmt.Errorf("read response body: %w", err) - } - return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody) - } - - id := resp.Header.Get("corro-query-id") - if id == "" { - resp.Body.Close() - return nil, errors.New("missing corro-query-id header in response") - } - - if skipRows { - return newSubscription(ctx, id, nil, resp.Body, nil), nil - } - - rows, err := newRows(ctx, resp.Body, false) - if err != nil { - resp.Body.Close() - return nil, fmt.Errorf("parse query response: %w", err) - } - return newSubscription(ctx, id, rows, rows.body, rows.decoder), nil -} - -func (c *APIClient) ResubscribeContext( - ctx context.Context, id string, skipRows bool, fromChange uint64, -) (*Subscription, error) { - // TODO - return nil, nil -} - -type ChangeType string - -var ( - ChangeTypeInsert ChangeType = "insert" - ChangeTypeUpdate ChangeType = "update" - ChangeTypeDelete ChangeType = "delete" -) - -type ChangeEvent struct { - Type ChangeType - RowID uint64 - Values []json.RawMessage - ChangeID uint64 -} - -func (ce *ChangeEvent) UnmarshalJSON(data []byte) error { - var raw []json.RawMessage - if err := json.Unmarshal(data, &raw); err != nil { - return fmt.Errorf("invalid change event: %w", err) - } - if len(raw) != 4 { - return fmt.Errorf("invalid change event: expected an array of 4 elements") - } - if err := json.Unmarshal(raw[0], &ce.Type); err != nil { - return fmt.Errorf("invalid change event type: %w", err) - } - if err := json.Unmarshal(raw[1], &ce.RowID); err != nil { - return fmt.Errorf("invalid change event row ID: %w", err) - } - if err := json.Unmarshal(raw[2], &ce.Values); err != nil { - return fmt.Errorf("invalid change event values: %w", err) - } - if err := json.Unmarshal(raw[3], &ce.ChangeID); err != nil { - return fmt.Errorf("invalid change event change ID: %w", err) - } - return nil -} - -func (ce *ChangeEvent) MarshalJSON() ([]byte, error) { - return json.Marshal([]any{ce.Type, ce.RowID, ce.Values, ce.ChangeID}) -} - -// Scan copies the column values in the change event into the values pointed at by dest. -// The number of values in dest must be the same as the number of columns in the change. -// Scan converts JSON-encoded column values to the provided Go types using [json.Unmarshal]. -func (ce *ChangeEvent) Scan(dest ...any) error { - if len(dest) != len(ce.Values) { - return fmt.Errorf("expected %d values, got %d", len(ce.Values), len(dest)) - } - - for i, v := range ce.Values { - if err := json.Unmarshal(v, dest[i]); err != nil { - return fmt.Errorf("unmarshal column value #%d: %w", i, err) - } - } - return nil -} - -// Subscription receives updates from the Corrosion database for a desired SQL query. -type Subscription struct { - ctx context.Context - cancel context.CancelFunc - - id string - rows *Rows - body io.ReadCloser - decoder *json.Decoder - changes chan *ChangeEvent - lastChangeID uint64 - err error -} - -func newSubscription( - ctx context.Context, id string, rows *Rows, body io.ReadCloser, decoder *json.Decoder, -) *Subscription { - ctx, cancel := context.WithCancel(ctx) - if decoder == nil { - decoder = json.NewDecoder(body) - } - return &Subscription{ - id: id, - rows: rows, - ctx: ctx, - cancel: cancel, - body: body, - decoder: decoder, - } -} - -// ID returns the subscription ID. -func (s *Subscription) ID() string { - return s.id -} - -// Rows returns the rows of the query or nil if skipRows was true when creating the subscription or if the subscription -// was created with [APIClient.ResubscribeContext]. -func (s *Subscription) Rows() *Rows { - return s.rows -} - -// Changes returns a channel that receives change events for the query. Changes are not available until all rows -// are consumed. The channel is closed when the context is done, or an error occurs while reading the changes, -// or when the subscription is closed explicitly. If it's closed due to an error, [Subscription.Err] will return -// the error. -func (s *Subscription) Changes() (<-chan *ChangeEvent, error) { - if s.changes != nil { - return s.changes, nil - } - - if s.rows != nil { - if s.rows.eoq == nil { - return nil, errors.New("changes are not available until all rows are consumed") - } - s.lastChangeID = *s.rows.eoq.ChangeID - } - s.changes = make(chan *ChangeEvent) - - go func() { - // Close the body when the context is done to unblock the decoder in the following goroutine. - <-s.ctx.Done() - s.body.Close() - }() - - go func() { - defer s.cancel() - defer close(s.changes) - - for { - select { - case <-s.ctx.Done(): - return - default: - } - - var e QueryEvent - if err := s.decoder.Decode(&e); err != nil { - // Do not report an error that occurred due to context cancellation. - if s.ctx.Err() == nil { - s.err = fmt.Errorf("decode query event: %w", err) - } - return - } - - if e.Error != nil { - s.err = fmt.Errorf("query error: %s", *e.Error) - return - } - if e.Change == nil { - s.err = fmt.Errorf("expected change event, got: %+v", e) - return - } - // If skipRows is true, the last change ID is unknown. - if s.lastChangeID != 0 && e.Change.ChangeID != s.lastChangeID+1 { - s.err = fmt.Errorf("missed a change: expected change ID %d, got %d", - s.lastChangeID+1, e.Change.ChangeID) - return - } - - s.lastChangeID = e.Change.ChangeID - select { - case s.changes <- e.Change: - case <-s.ctx.Done(): - return - } - } - }() - - return s.changes, nil -} - -// Err returns the error, if any, that was encountered during fetching changes. -// Err may be called after an explicit or implicit [Subscription.Close]. -func (s *Subscription) Err() error { - return s.err -} - -func (s *Subscription) Close() error { - s.cancel() - return s.body.Close() -} diff --git a/internal/corrosion/subscribe.go b/internal/corrosion/subscribe.go new file mode 100644 index 00000000..ae9ec45e --- /dev/null +++ b/internal/corrosion/subscribe.go @@ -0,0 +1,253 @@ +package corrosion + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" +) + +// SubscribeContext creates a subscription to receive updates for a desired SQL query. If skipRows is false, +// Subscription.Rows must be consumed before Subscription.Changes can be called. If skipRows is true, Subscription.Rows +// will be nil. +func (c *APIClient) SubscribeContext( + ctx context.Context, query string, args []any, skipRows bool, +) (*Subscription, error) { + statement := Statement{ + Query: query, + Params: args, + } + body, err := json.Marshal(statement) + if err != nil { + return nil, fmt.Errorf("marshal query: %w", err) + } + + subURL := c.baseURL.JoinPath("/v1/subscriptions") + if skipRows { + q := subURL.Query() + q.Set("skip_rows", "true") + subURL.RawQuery = q.Encode() + } + + req, err := http.NewRequestWithContext(ctx, "POST", subURL.String(), bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + + resp, err := c.client.Do(req) + if err != nil { + return nil, fmt.Errorf("send request: %w", err) + } + + if resp.StatusCode != http.StatusOK { + respBody, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil { + return nil, fmt.Errorf("read response body: %w", err) + } + return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody) + } + + id := resp.Header.Get("corro-query-id") + if id == "" { + resp.Body.Close() + return nil, errors.New("missing corro-query-id header in response") + } + + if skipRows { + return newSubscription(ctx, id, nil, resp.Body, nil), nil + } + + rows, err := newRows(ctx, resp.Body, false) + if err != nil { + resp.Body.Close() + return nil, fmt.Errorf("parse query response: %w", err) + } + return newSubscription(ctx, id, rows, rows.body, rows.decoder), nil +} + +type ChangeType string + +var ( + ChangeTypeInsert ChangeType = "insert" + ChangeTypeUpdate ChangeType = "update" + ChangeTypeDelete ChangeType = "delete" +) + +type ChangeEvent struct { + Type ChangeType + RowID uint64 + Values []json.RawMessage + ChangeID uint64 +} + +func (ce *ChangeEvent) UnmarshalJSON(data []byte) error { + var raw []json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return fmt.Errorf("invalid change event: %w", err) + } + if len(raw) != 4 { + return fmt.Errorf("invalid change event: expected an array of 4 elements") + } + if err := json.Unmarshal(raw[0], &ce.Type); err != nil { + return fmt.Errorf("invalid change event type: %w", err) + } + if err := json.Unmarshal(raw[1], &ce.RowID); err != nil { + return fmt.Errorf("invalid change event row ID: %w", err) + } + if err := json.Unmarshal(raw[2], &ce.Values); err != nil { + return fmt.Errorf("invalid change event values: %w", err) + } + if err := json.Unmarshal(raw[3], &ce.ChangeID); err != nil { + return fmt.Errorf("invalid change event change ID: %w", err) + } + return nil +} + +func (ce *ChangeEvent) MarshalJSON() ([]byte, error) { + return json.Marshal([]any{ce.Type, ce.RowID, ce.Values, ce.ChangeID}) +} + +// Scan copies the column values in the change event into the values pointed at by dest. +// The number of values in dest must be the same as the number of columns in the change. +// Scan converts JSON-encoded column values to the provided Go types using [json.Unmarshal]. +func (ce *ChangeEvent) Scan(dest ...any) error { + if len(dest) != len(ce.Values) { + return fmt.Errorf("expected %d values, got %d", len(ce.Values), len(dest)) + } + + for i, v := range ce.Values { + if err := json.Unmarshal(v, dest[i]); err != nil { + return fmt.Errorf("unmarshal column value #%d: %w", i, err) + } + } + return nil +} + +// Subscription receives updates from the Corrosion database for a desired SQL query. +type Subscription struct { + ctx context.Context + cancel context.CancelFunc + + id string + rows *Rows + body io.ReadCloser + decoder *json.Decoder + changes chan *ChangeEvent + lastChangeID uint64 + err error +} + +func newSubscription( + ctx context.Context, id string, rows *Rows, body io.ReadCloser, decoder *json.Decoder, +) *Subscription { + ctx, cancel := context.WithCancel(ctx) + if decoder == nil { + decoder = json.NewDecoder(body) + } + return &Subscription{ + id: id, + rows: rows, + ctx: ctx, + cancel: cancel, + body: body, + decoder: decoder, + } +} + +// ID returns the subscription ID. +func (s *Subscription) ID() string { + return s.id +} + +// Rows returns the rows of the query or nil if skipRows was true when creating the subscription or if the subscription +// was created with [APIClient.ResubscribeContext]. +func (s *Subscription) Rows() *Rows { + return s.rows +} + +// Changes returns a channel that receives change events for the query. Changes are not available until all rows +// are consumed. The channel is closed when the context is done, or an error occurs while reading the changes, +// or when the subscription is closed explicitly. If it's closed due to an error, [Subscription.Err] will return +// the error. +func (s *Subscription) Changes() (<-chan *ChangeEvent, error) { + if s.changes != nil { + return s.changes, nil + } + + if s.rows != nil { + if s.rows.eoq == nil { + return nil, errors.New("changes are not available until all rows are consumed") + } + s.lastChangeID = *s.rows.eoq.ChangeID + } + s.changes = make(chan *ChangeEvent) + + go func() { + // Close the body when the context is done to unblock the decoder in the following goroutine. + <-s.ctx.Done() + s.body.Close() + }() + + go func() { + defer s.cancel() + defer close(s.changes) + + for { + select { + case <-s.ctx.Done(): + return + default: + } + + var e QueryEvent + if err := s.decoder.Decode(&e); err != nil { + // Do not report an error that occurred due to context cancellation. + if s.ctx.Err() == nil { + s.err = fmt.Errorf("decode query event: %w", err) + } + return + } + + if e.Error != nil { + s.err = fmt.Errorf("query error: %s", *e.Error) + return + } + if e.Change == nil { + s.err = fmt.Errorf("expected change event, got: %+v", e) + return + } + // If skipRows is true, the last change ID is unknown. + if s.lastChangeID != 0 && e.Change.ChangeID != s.lastChangeID+1 { + s.err = fmt.Errorf("missed a change: expected change ID %d, got %d", + s.lastChangeID+1, e.Change.ChangeID) + return + } + + s.lastChangeID = e.Change.ChangeID + select { + case s.changes <- e.Change: + case <-s.ctx.Done(): + return + } + } + }() + + return s.changes, nil +} + +// Err returns the error, if any, that was encountered during fetching changes. +// Err may be called after an explicit or implicit [Subscription.Close]. +func (s *Subscription) Err() error { + return s.err +} + +func (s *Subscription) Close() error { + s.cancel() + return s.body.Close() +}