hatchet/pkg/repository/v1/sqlcv1/queue.sql.go

// Code generated by sqlc. DO NOT EDIT.
// versions:
//   sqlc v1.24.0
// source: queue.sql

package sqlcv1

import (
	"context"

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

const bulkQueueItems = `-- name: BulkQueueItems :many
WITH locked_qis AS (
    SELECT
        id
    FROM
        v1_queue_item
    WHERE
        id = ANY($1::bigint[])
    ORDER BY
        id ASC
    FOR UPDATE
)
DELETE FROM
    v1_queue_item
WHERE
    id = ANY($1::bigint[])
RETURNING
    id
`

func (q *Queries) BulkQueueItems(ctx context.Context, db DBTX, ids []int64) ([]int64, error) {
	rows, err := db.Query(ctx, bulkQueueItems, ids)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []int64
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		items = append(items, id)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const deleteTasksFromQueue = `-- name: DeleteTasksFromQueue :exec
WITH input AS (
    SELECT
        task_id, retry_count
    FROM
        (
            SELECT
                unnest($1::bigint[]) AS task_id,
                unnest($2::integer[]) AS retry_count
        ) AS subquery
), locked_qis AS (
    SELECT
        id
    FROM
        v1_queue_item
    WHERE
        (task_id, retry_count) IN (SELECT task_id, retry_count FROM input)
    ORDER BY
        id ASC
    FOR UPDATE
)
DELETE FROM
    v1_queue_item
WHERE
    id = ANY(SELECT id FROM locked_qis)
`

type DeleteTasksFromQueueParams struct {
	Taskids     []int64 `json:"taskids"`
	Retrycounts []int32 `json:"retrycounts"`
}

func (q *Queries) DeleteTasksFromQueue(ctx context.Context, db DBTX, arg DeleteTasksFromQueueParams) error {
	_, err := db.Exec(ctx, deleteTasksFromQueue, arg.Taskids, arg.Retrycounts)
	return err
}

const getDesiredLabels = `-- name: GetDesiredLabels :many
SELECT
    "key",
    "strValue",
    "intValue",
    "required",
    "weight",
    "comparator",
    "stepId"
FROM
    "StepDesiredWorkerLabel"
WHERE
    "stepId" = ANY($1::uuid[])
`

type GetDesiredLabelsRow struct {
	Key        string                `json:"key"`
	StrValue   pgtype.Text           `json:"strValue"`
	IntValue   pgtype.Int4           `json:"intValue"`
	Required   bool                  `json:"required"`
	Weight     int32                 `json:"weight"`
	Comparator WorkerLabelComparator `json:"comparator"`
	StepId     pgtype.UUID           `json:"stepId"`
}

func (q *Queries) GetDesiredLabels(ctx context.Context, db DBTX, stepids []pgtype.UUID) ([]*GetDesiredLabelsRow, error) {
	rows, err := db.Query(ctx, getDesiredLabels, stepids)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*GetDesiredLabelsRow
	for rows.Next() {
		var i GetDesiredLabelsRow
		if err := rows.Scan(
			&i.Key,
			&i.StrValue,
			&i.IntValue,
			&i.Required,
			&i.Weight,
			&i.Comparator,
			&i.StepId,
		); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const getMinUnprocessedQueueItemId = `-- name: GetMinUnprocessedQueueItemId :one
WITH priority_1 AS (
    SELECT
        id
    FROM
        v1_queue_item
    WHERE
        tenant_id = $1::uuid
        AND queue = $2::text
        AND priority = 1
    ORDER BY
        id ASC
    LIMIT 1
),
priority_2 AS (
    SELECT
        id
    FROM
        v1_queue_item
    WHERE
        tenant_id = $1::uuid
        AND queue = $2::text
        AND priority = 2
    ORDER BY
        id ASC
    LIMIT 1
),
priority_3 AS (
    SELECT
        id
    FROM
        v1_queue_item
    WHERE
        tenant_id = $1::uuid
        AND queue = $2::text
        AND priority = 3
    ORDER BY
        id ASC
    LIMIT 1
),
priority_4 AS (
    SELECT
        id
    FROM
        v1_queue_item
    WHERE
        tenant_id = $1::uuid
        AND queue = $2::text
        AND priority = 4
    ORDER BY
        id ASC
    LIMIT 1
)
SELECT
    COALESCE(MIN(id), 0)::bigint AS "minId"
FROM (
    SELECT id FROM priority_1
    UNION ALL
    SELECT id FROM priority_2
    UNION ALL
    SELECT id FROM priority_3
    UNION ALL
    SELECT id FROM priority_4
) AS combined_priorities
`

type GetMinUnprocessedQueueItemIdParams struct {
	Tenantid pgtype.UUID `json:"tenantid"`
	Queue    string      `json:"queue"`
}

func (q *Queries) GetMinUnprocessedQueueItemId(ctx context.Context, db DBTX, arg GetMinUnprocessedQueueItemIdParams) (int64, error) {
	row := db.QueryRow(ctx, getMinUnprocessedQueueItemId, arg.Tenantid, arg.Queue)
	var minId int64
	err := row.Scan(&minId)
	return minId, err
}

const getQueuedCounts = `-- name: GetQueuedCounts :many
SELECT
    queue,
    COUNT(*) AS count
FROM
    v1_queue_item qi
WHERE
    qi.tenant_id = $1::uuid
GROUP BY
    qi.queue
`

type GetQueuedCountsRow struct {
	Queue string `json:"queue"`
	Count int64  `json:"count"`
}

func (q *Queries) GetQueuedCounts(ctx context.Context, db DBTX, tenantid pgtype.UUID) ([]*GetQueuedCountsRow, error) {
	rows, err := db.Query(ctx, getQueuedCounts, tenantid)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*GetQueuedCountsRow
	for rows.Next() {
		var i GetQueuedCountsRow
		if err := rows.Scan(&i.Queue, &i.Count); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const listActionsForWorkers = `-- name: ListActionsForWorkers :many
SELECT
    w."id" as "workerId",
    a."actionId"
FROM
    "Worker" w
LEFT JOIN
    "_ActionToWorker" atw ON w."id" = atw."B"
LEFT JOIN
    "Action" a ON atw."A" = a."id"
WHERE
    w."tenantId" = $1::uuid
    AND w."id" = ANY($2::uuid[])
    AND w."dispatcherId" IS NOT NULL
    AND w."lastHeartbeatAt" > NOW() - INTERVAL '5 seconds'
    AND w."isActive" = true
    AND w."isPaused" = false
`

type ListActionsForWorkersParams struct {
	Tenantid  pgtype.UUID   `json:"tenantid"`
	Workerids []pgtype.UUID `json:"workerids"`
}

type ListActionsForWorkersRow struct {
	WorkerId pgtype.UUID `json:"workerId"`
	ActionId pgtype.Text `json:"actionId"`
}

func (q *Queries) ListActionsForWorkers(ctx context.Context, db DBTX, arg ListActionsForWorkersParams) ([]*ListActionsForWorkersRow, error) {
	rows, err := db.Query(ctx, listActionsForWorkers, arg.Tenantid, arg.Workerids)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*ListActionsForWorkersRow
	for rows.Next() {
		var i ListActionsForWorkersRow
		if err := rows.Scan(&i.WorkerId, &i.ActionId); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const listAvailableSlotsForWorkers = `-- name: ListAvailableSlotsForWorkers :many
WITH worker_max_runs AS (
    SELECT
        "id",
        "maxRuns"
    FROM
        "Worker"
    WHERE
        "tenantId" = $1::uuid
        AND "id" = ANY($2::uuid[])
), worker_filled_slots AS (
    SELECT
        worker_id,
        COUNT(task_id) AS "filledSlots"
    FROM
        v1_task_runtime
    WHERE
        tenant_id = $1::uuid
        AND worker_id = ANY($2::uuid[])
    GROUP BY
        worker_id
)
SELECT
    wmr."id",
    wmr."maxRuns" - COALESCE(wfs."filledSlots", 0) AS "availableSlots"
FROM
    worker_max_runs wmr
LEFT JOIN
    worker_filled_slots wfs ON wmr."id" = wfs.worker_id
`

type ListAvailableSlotsForWorkersParams struct {
	Tenantid  pgtype.UUID   `json:"tenantid"`
	Workerids []pgtype.UUID `json:"workerids"`
}

type ListAvailableSlotsForWorkersRow struct {
	ID             pgtype.UUID `json:"id"`
	AvailableSlots int32       `json:"availableSlots"`
}

// subtract the filled slots from the max runs to get the available slots
func (q *Queries) ListAvailableSlotsForWorkers(ctx context.Context, db DBTX, arg ListAvailableSlotsForWorkersParams) ([]*ListAvailableSlotsForWorkersRow, error) {
	rows, err := db.Query(ctx, listAvailableSlotsForWorkers, arg.Tenantid, arg.Workerids)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*ListAvailableSlotsForWorkersRow
	for rows.Next() {
		var i ListAvailableSlotsForWorkersRow
		if err := rows.Scan(&i.ID, &i.AvailableSlots); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const listQueueItemsForQueue = `-- name: ListQueueItemsForQueue :many
SELECT
    id, tenant_id, queue, task_id, task_inserted_at, external_id, action_id, step_id, workflow_id, workflow_run_id, schedule_timeout_at, step_timeout, priority, sticky, desired_worker_id, retry_count
FROM
    v1_queue_item qi
WHERE
    qi.tenant_id = $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)
`

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

func (q *Queries) ListQueueItemsForQueue(ctx context.Context, db DBTX, arg ListQueueItemsForQueueParams) ([]*V1QueueItem, error) {
	rows, err := db.Query(ctx, listQueueItemsForQueue,
		arg.Tenantid,
		arg.Queue,
		arg.GtId,
		arg.Limit,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*V1QueueItem
	for rows.Next() {
		var i V1QueueItem
		if err := rows.Scan(
			&i.ID,
			&i.TenantID,
			&i.Queue,
			&i.TaskID,
			&i.TaskInsertedAt,
			&i.ExternalID,
			&i.ActionID,
			&i.StepID,
			&i.WorkflowID,
			&i.WorkflowRunID,
			&i.ScheduleTimeoutAt,
			&i.StepTimeout,
			&i.Priority,
			&i.Sticky,
			&i.DesiredWorkerID,
			&i.RetryCount,
		); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const listQueues = `-- name: ListQueues :many
SELECT
    tenant_id, name, last_active
FROM
    v1_queue
WHERE
    tenant_id = $1::uuid
    AND last_active > NOW() - INTERVAL '1 day'
`

func (q *Queries) ListQueues(ctx context.Context, db DBTX, tenantid pgtype.UUID) ([]*V1Queue, error) {
	rows, err := db.Query(ctx, listQueues, tenantid)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*V1Queue
	for rows.Next() {
		var i V1Queue
		if err := rows.Scan(&i.TenantID, &i.Name, &i.LastActive); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const updateTasksToAssigned = `-- name: UpdateTasksToAssigned :many
WITH input AS (
    SELECT
        id,
        worker_id
    FROM
        (
            SELECT
                unnest($1::bigint[]) AS id,
                unnest($2::uuid[]) AS worker_id
        ) AS subquery
    ORDER BY id
), updated_tasks AS (
    SELECT
        t.id,
        t.inserted_at,
        t.retry_count,
        input.worker_id,
        t.tenant_id,
        CURRENT_TIMESTAMP + convert_duration_to_interval(t.step_timeout) AS timeout_at
    FROM
        input
    JOIN
        v1_task t ON t.id = input.id
    ORDER BY t.id
), assigned_tasks AS (
    INSERT INTO v1_task_runtime (
        task_id,
        task_inserted_at,
        retry_count,
        worker_id,
        tenant_id,
        timeout_at
    )
    SELECT
        t.id,
        t.inserted_at,
        t.retry_count,
        t.worker_id,
        $3::uuid,
        t.timeout_at
    FROM
        updated_tasks t
    ON CONFLICT (task_id, task_inserted_at, retry_count) DO NOTHING
    -- only return the task ids that were successfully assigned
    RETURNING task_id, worker_id
)
SELECT
    asr.task_id,
    asr.worker_id
FROM
    assigned_tasks asr
`

type UpdateTasksToAssignedParams struct {
	Taskids   []int64       `json:"taskids"`
	Workerids []pgtype.UUID `json:"workerids"`
	Tenantid  pgtype.UUID   `json:"tenantid"`
}

type UpdateTasksToAssignedRow struct {
	TaskID   int64       `json:"task_id"`
	WorkerID pgtype.UUID `json:"worker_id"`
}

func (q *Queries) UpdateTasksToAssigned(ctx context.Context, db DBTX, arg UpdateTasksToAssignedParams) ([]*UpdateTasksToAssignedRow, error) {
	rows, err := db.Query(ctx, updateTasksToAssigned, arg.Taskids, arg.Workerids, arg.Tenantid)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []*UpdateTasksToAssignedRow
	for rows.Next() {
		var i UpdateTasksToAssignedRow
		if err := rows.Scan(&i.TaskID, &i.WorkerID); err != nil {
			return nil, err
		}
		items = append(items, &i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const upsertQueues = `-- name: UpsertQueues :exec
WITH ordered_names AS (
    SELECT unnest($2::text[]) AS name
    ORDER BY name
)
INSERT INTO
    v1_queue (
        tenant_id,
        name,
        last_active
    )
SELECT
    $1,
    name,
    NOW()
FROM ordered_names
ON CONFLICT (tenant_id, name) DO UPDATE
SET
    last_active = NOW()
`

type UpsertQueuesParams struct {
	TenantID pgtype.UUID `json:"tenant_id"`
	Names    []string    `json:"names"`
}

func (q *Queries) UpsertQueues(ctx context.Context, db DBTX, arg UpsertQueuesParams) error {
	_, err := db.Exec(ctx, upsertQueues, arg.TenantID, arg.Names)
	return err
}