diff --git a/internal/corrosion/client.go b/internal/corrosion/client.go index c9a03174..b8e75d5f 100644 --- a/internal/corrosion/client.go +++ b/internal/corrosion/client.go @@ -39,7 +39,8 @@ func NewAPIClient(addr netip.AddrPort) (*APIClient, error) { return &APIClient{ baseURL: baseURL, client: &http.Client{ - Timeout: http2Timeout, + // TODO: use timeout for non-subscription requests? + //Timeout: http2Timeout, Transport: &RetryRoundTripper{ Base: &http2.Transport{ AllowHTTP: true, @@ -181,11 +182,10 @@ func (c *APIClient) ExecMultiContext(ctx context.Context, statements ...Statemen } type QueryEvent struct { - Columns []string `json:"columns"` - Row *RowEvent `json:"row"` - EOQ *EndOfQuery `json:"eoq"` - // TODO: implement event type Change to support subscriptions. - //Change []any `json:"change"` + Columns []string `json:"columns"` + Row *RowEvent `json:"row"` + EOQ *EndOfQuery `json:"eoq"` + Change *ChangeEvent `json:"change"` // Error is a server-side error that occurred during query execution. It's considered fatal for the client // as it cannot be recovered from server-side. Error *string `json:"error"` @@ -256,7 +256,7 @@ func (c *APIClient) QueryContext(ctx context.Context, query string, args ...any) return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody) } - rows, err := newRows(ctx, resp.Body) + rows, err := newRows(ctx, resp.Body, true) if err != nil { resp.Body.Close() return nil, fmt.Errorf("parse query response: %w", err) @@ -267,17 +267,18 @@ func (c *APIClient) QueryContext(ctx context.Context, query string, args ...any) // Rows is the result of a query. Its cursor starts before the first row of the result set. // Use [Rows.Next] to advance from row to row. type Rows struct { - ctx context.Context - body io.ReadCloser - decoder *json.Decoder + ctx context.Context + body io.ReadCloser + decoder *json.Decoder + eoq *EndOfQuery + closeOnEOQ bool columns []string row RowEvent - time float64 err error } -func newRows(ctx context.Context, body io.ReadCloser) (*Rows, error) { +func newRows(ctx context.Context, body io.ReadCloser, closeOnEOQ bool) (*Rows, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -294,10 +295,11 @@ func newRows(ctx context.Context, body io.ReadCloser) (*Rows, error) { } return &Rows{ - ctx: ctx, - body: body, - decoder: decoder, - columns: e.Columns, + ctx: ctx, + body: body, + decoder: decoder, + closeOnEOQ: closeOnEOQ, + columns: e.Columns, }, nil } @@ -343,8 +345,11 @@ func (rs *Rows) Next() bool { return true } if e.EOQ != nil { - rs.time = e.EOQ.Time - _ = rs.Close() + rs.eoq = e.EOQ + // Rows could be used as part of a subscription, so don't close the body if so. + if rs.closeOnEOQ { + _ = rs.Close() + } return false } @@ -381,13 +386,13 @@ func (rs *Rows) Scan(dest ...any) error { // Time returns the time taken to execute the query in seconds. It's only available after all rows have been consumed. // It doesn't include the time to send the query, receive the response, or iterate over the rows. func (rs *Rows) Time() (float64, error) { - if rs.time == 0 { + if rs.eoq == nil { if rs.Err() != nil { return 0, fmt.Errorf("time is not available: %w", rs.Err()) } return 0, errors.New("time is not available until all rows are consumed") } - return rs.time, nil + return rs.eoq.Time, nil } // Close closes the [Rows], preventing further enumeration. If [Rows.Next] is called and returns false, @@ -396,3 +401,242 @@ func (rs *Rows) Time() (float64, error) { func (rs *Rows) Close() error { return rs.body.Close() } + +// 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 { + ID string + Rows *Rows + + ctx context.Context + cancel context.CancelFunc + 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, + } +} + +// 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 { + if s.ctx.Err() != nil { + s.err = s.ctx.Err() + } else { + 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() +}