hatchet/pkg/repository/prisma/dbsqlc/batch.go

// Code generated by sqlc. DO NOT EDIT.
// versions:
//   sqlc v1.24.0
// source: batch.go

package dbsqlc

import (
	"context"
	"errors"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
)

var (
	ErrBatchAlreadyClosed = errors.New("batch already closed")
)

const listQueueItems = `-- name: ListQueueItems :batchmany
SELECT
    id, "stepRunId", "stepId", "actionId", "scheduleTimeoutAt", "stepTimeout", priority, "isQueued", "tenantId", queue, sticky, "desiredWorkerId"
FROM
    "QueueItem" qi
WHERE
    qi."isQueued" = true
    AND qi."tenantId" = $1::uuid
    AND qi."queue" = $2::text
    AND (
        $3::bigint IS NULL OR
        qi."id" >= $3::bigint
    )
    -- Added to ensure that the index is used
    AND qi."priority" >= 1 AND qi."priority" <= 4
ORDER BY
    qi."priority" DESC,
    qi."id" ASC
LIMIT
    COALESCE($4::integer, 100)
FOR UPDATE SKIP LOCKED
`

type ListQueueItemsBatchResults struct {
	br     pgx.BatchResults
	tot    int
	closed bool
}

type ListQueueItemsParams struct {
	Tenantid pgtype.UUID `json:"tenantid"`
	Queue    string      `json:"queue"`
	GtId     pgtype.Int8 `json:"gtId"`
	Limit    pgtype.Int4 `json:"limit"`
}

func (q *Queries) ListQueueItems(ctx context.Context, db DBTX, arg []ListQueueItemsParams) *ListQueueItemsBatchResults {
	batch := &pgx.Batch{}
	for _, a := range arg {
		vals := []interface{}{
			a.Tenantid,
			a.Queue,
			a.GtId,
			a.Limit,
		}
		batch.Queue(listQueueItems, vals...)
	}
	br := db.SendBatch(ctx, batch)
	return &ListQueueItemsBatchResults{br, len(arg), false}
}

func (b *ListQueueItemsBatchResults) Query(f func(int, []*QueueItem, error)) {
	defer b.br.Close()
	for t := 0; t < b.tot; t++ {
		var items []*QueueItem
		if b.closed {
			if f != nil {
				f(t, items, ErrBatchAlreadyClosed)
			}
			continue
		}
		err := func() error {
			rows, err := b.br.Query()
			if err != nil {
				return err
			}
			defer rows.Close()
			for rows.Next() {
				var i QueueItem
				if err := rows.Scan(
					&i.ID,
					&i.StepRunId,
					&i.StepId,
					&i.ActionId,
					&i.ScheduleTimeoutAt,
					&i.StepTimeout,
					&i.Priority,
					&i.IsQueued,
					&i.TenantId,
					&i.Queue,
					&i.Sticky,
					&i.DesiredWorkerId,
				); err != nil {
					return err
				}
				items = append(items, &i)
			}
			return rows.Err()
		}()
		if f != nil {
			f(t, items, err)
		}
	}
}

func (b *ListQueueItemsBatchResults) Close() error {
	b.closed = true
	return b.br.Close()
}