split corrosion client subscription in a separate file

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

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