package sqldb

import (

"context"

"fmt"

)

const (

// maxSQLiteBatchSize is the maximum number of items that can be

// included in a batch query IN clause for SQLite. This was determined

// using the TestSQLSliceQueries test.

maxSQLiteBatchSize = 32766

// maxPostgresBatchSize is the maximum number of items that can be

// included in a batch query IN clause for Postgres. This was determined

// using the TestSQLSliceQueries test.

maxPostgresBatchSize = 65535

// defaultSQLitePageSize is the default page size for SQLite queries.

defaultSQLitePageSize = 100

// defaultPostgresPageSize is the default page size for Postgres

// queries.

defaultPostgresPageSize = 10500

// defaultSQLiteBatchSize is the default batch size for SQLite queries.

defaultSQLiteBatchSize = 250

// defaultPostgresBatchSize is the default batch size for Postgres

// queries.

defaultPostgresBatchSize = 5000

)

// QueryConfig holds configuration values for SQL queries.

//

//nolint:ll

type QueryConfig struct {

// MaxBatchSize is the maximum number of items included in a batch

// query IN clauses list.

MaxBatchSize uint32 `long:"max-batch-size" description:"The maximum number of items to include in a batch query IN clause. This is used for queries that fetch results based on a list of identifiers."`

// MaxPageSize is the maximum number of items returned in a single page

// of results. This is used for paginated queries.

MaxPageSize uint32 `long:"max-page-size" description:"The maximum number of items to return in a single page of results. This is used for paginated queries."`

}

// Validate checks that the QueryConfig values are valid.

func (c *QueryConfig) Validate(sqlite bool) error {

if c.MaxBatchSize <= 0 {

return fmt.Errorf("max batch size must be greater than "+

"zero, got %d", c.MaxBatchSize)

}

if c.MaxPageSize <= 0 {

return fmt.Errorf("max page size must be greater than "+

"zero, got %d", c.MaxPageSize)

}

if sqlite {

if c.MaxBatchSize > maxSQLiteBatchSize {

return fmt.Errorf("max batch size for SQLite cannot "+

"exceed %d, got %d", maxSQLiteBatchSize,

c.MaxBatchSize)

}

} else {

if c.MaxBatchSize > maxPostgresBatchSize {

return fmt.Errorf("max batch size for Postgres cannot "+

"exceed %d, got %d", maxPostgresBatchSize,

c.MaxBatchSize)

}

}

return nil

}

// DefaultSQLiteConfig returns a default configuration for SQL queries to a

// SQLite backend.

func DefaultSQLiteConfig() *QueryConfig {

return &QueryConfig{

MaxBatchSize: defaultSQLiteBatchSize,

MaxPageSize: defaultSQLitePageSize,

}

}

// DefaultPostgresConfig returns a default configuration for SQL queries to a

// Postgres backend.

func DefaultPostgresConfig() *QueryConfig {

return &QueryConfig{

MaxBatchSize: defaultPostgresBatchSize,

MaxPageSize: defaultPostgresPageSize,

}

}

// 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 += int(cfg.MaxBatchSize) {

// Calculate the end index for this page.

end := i + int(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, int32(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

}

// CollectAndBatchDataQueryFunc represents a function that batch loads

// additional data for collected identifiers, returning the batch data that

// applies to all items.

type CollectAndBatchDataQueryFunc[ID any, BatchData any] func(context.Context,

[]ID) (BatchData, error)

// ItemWithBatchDataProcessFunc represents a function that processes individual

// items along with shared batch data.

type ItemWithBatchDataProcessFunc[T any, BatchData any] func(context.Context,

T, BatchData) error

// CollectFunc represents a function that extracts an identifier from a

// paginated item.

type CollectFunc[T any, ID any] func(T) (ID, error)

// ExecuteCollectAndBatchWithSharedDataQuery implements a page-by-page

// processing pattern where each page is immediately processed with batch-loaded

// data before moving to the next page.

//

// It:

// 1. Fetches a page of items using cursor-based pagination

// 2. Collects identifiers from that page and batch loads shared data

// 3. Processes each item in the page with the shared batch data

// 4. Moves to the next page and repeats

//

// Parameters:

// - initialCursor: starting cursor for pagination

// - pageQueryFunc: fetches a page of items

// - extractPageCursor: extracts cursor from paginated item for next page

// - collectFunc: extracts identifier from paginated item

// - batchDataFunc: batch loads shared data from collected IDs for one page

// - processItem: processes each item with the shared batch data

func ExecuteCollectAndBatchWithSharedDataQuery[C any, T any, I any, D any](

ctx context.Context, cfg *QueryConfig, initialCursor C,

pageQueryFunc PagedQueryFunc[C, T],

extractPageCursor CursorExtractFunc[T, C],

collectFunc CollectFunc[T, I],

batchDataFunc CollectAndBatchDataQueryFunc[I, D],

processItem ItemWithBatchDataProcessFunc[T, D]) error {

cursor := initialCursor

for {

// Step 1: Fetch the next page of items.

items, err := pageQueryFunc(ctx, cursor, int32(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

}

// Step 2: Collect identifiers from this page and batch load

// data.

pageIDs := make([]I, len(items))

for i, item := range items {

pageIDs[i], err = collectFunc(item)

if err != nil {

return fmt.Errorf("failed to collect "+

"identifier from item: %w", err)

}

}

// Batch load shared data for this page.

batchData, err := batchDataFunc(ctx, pageIDs)

if err != nil {

return fmt.Errorf("failed to load batch data for "+

"page: %w", err)

}

// Step 3: Process each item in this page with the shared batch

// data.

for _, item := range items {

err := processItem(ctx, item, batchData)

if err != nil {

return fmt.Errorf("failed to process item "+

"with batch data: %w", err)

}

// Update cursor for next page.

cursor = extractPageCursor(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

}