package sqldb import ( "context" "fmt" ) // QueryConfig holds configuration values for SQL queries. type QueryConfig struct { // MaxBatchSize is the maximum number of items included in a batch // query IN clauses list. MaxBatchSize int // MaxPageSize is the maximum number of items returned in a single page // of results. This is used for paginated queries. MaxPageSize int32 } // DefaultQueryConfig returns a default configuration for SQL queries. // // TODO(elle): make configurable & have different defaults for SQLite and // Postgres. func DefaultQueryConfig() *QueryConfig { return &QueryConfig{ MaxBatchSize: 250, MaxPageSize: 10000, } } // BatchQueryFunc represents a function that takes a batch of converted items // and returns results. type BatchQueryFunc[T any, R any] func(context.Context, []T) ([]R, error) // ItemCallbackFunc represents a function that processes individual results. type ItemCallbackFunc[R any] func(context.Context, R) error // ConvertFunc represents a function that converts from input type to query type // for the batch query. type ConvertFunc[I any, T any] func(I) T // ExecuteBatchQuery executes a query in batches over a slice of input items. // It converts the input items to a query type using the provided convertFunc, // executes the query in batches using the provided queryFunc, and applies // the callback to each result. This is useful for queries using the // "WHERE x IN []slice" pattern. It takes that slice, splits it into batches of // size MaxBatchSize, and executes the query for each batch. // // NOTE: it is the caller's responsibility to ensure that the expected return // results are unique across all pages. Meaning that if the input items are // split up, a result that is returned in one page should not be expected to // be returned in another page. func ExecuteBatchQuery[I any, T any, R any](ctx context.Context, cfg *QueryConfig, inputItems []I, convertFunc ConvertFunc[I, T], queryFunc BatchQueryFunc[T, R], callback ItemCallbackFunc[R]) error { if len(inputItems) == 0 { return nil } // Process items in pages. for i := 0; i < len(inputItems); i += cfg.MaxBatchSize { // Calculate the end index for this page. end := i + cfg.MaxBatchSize if end > len(inputItems) { end = len(inputItems) } // Get the page slice of input items. inputPage := inputItems[i:end] // Convert only the items needed for this page. convertedPage := make([]T, len(inputPage)) for j, inputItem := range inputPage { convertedPage[j] = convertFunc(inputItem) } // Execute the query for this page. results, err := queryFunc(ctx, convertedPage) if err != nil { return fmt.Errorf("query failed for page "+ "starting at %d: %w", i, err) } // Apply the callback to each result. for _, result := range results { if err := callback(ctx, result); err != nil { return fmt.Errorf("callback failed for "+ "result: %w", err) } } } return nil } // PagedQueryFunc represents a function that fetches a page of results using a // cursor. It returns the fetched items and should return an empty slice when no // more results. type PagedQueryFunc[C any, T any] func(context.Context, C, int32) ([]T, error) // CursorExtractFunc represents a function that extracts the cursor value from // an item. This cursor will be used for the next page fetch. type CursorExtractFunc[T any, C any] func(T) C // ItemProcessFunc represents a function that processes individual items. type ItemProcessFunc[T any] func(context.Context, T) error // ExecutePaginatedQuery executes a cursor-based paginated query. It continues // fetching pages until no more results are returned, processing each item with // the provided callback. // // Parameters: // - initialCursor: the starting cursor value (e.g., 0, -1, "", etc.). // - queryFunc: function that fetches a page given cursor and limit. // - extractCursor: function that extracts cursor from an item for next page. // - processItem: function that processes each individual item. // // NOTE: it is the caller's responsibility to "undo" any processing done on // items if the query fails on a later page. func ExecutePaginatedQuery[C any, T any](ctx context.Context, cfg *QueryConfig, initialCursor C, queryFunc PagedQueryFunc[C, T], extractCursor CursorExtractFunc[T, C], processItem ItemProcessFunc[T]) error { cursor := initialCursor for { // Fetch the next page. items, err := queryFunc(ctx, cursor, cfg.MaxPageSize) if err != nil { return fmt.Errorf("failed to fetch page with "+ "cursor %v: %w", cursor, err) } // If no items returned, we're done. if len(items) == 0 { break } // Process each item in the page. for _, item := range items { if err := processItem(ctx, item); err != nil { return fmt.Errorf("failed to process item: %w", err) } // Update cursor for next iteration. cursor = extractCursor(item) } // If the number of items is less than the max page size, // we assume there are no more items to fetch. if len(items) < int(cfg.MaxPageSize) { break } } return nil }