workers.sql

-- name: ListManyWorkerLabels :many
SELECT
    "id",
    "key",
    "intValue",
    "strValue",
    "createdAt",
    "updatedAt",
    "workerId"
FROM "WorkerLabel" wl
WHERE wl."workerId" = ANY(@workerIds::uuid[]);

-- name: ListWorkerSlotConfigs :many
SELECT
    worker_id,
    slot_type,
    max_units
FROM
    v1_worker_slot_config
WHERE
    tenant_id = @tenantId::uuid
    AND worker_id = ANY(@workerIds::uuid[]);

-- name: CreateWorkerSlotConfigs :exec
INSERT INTO v1_worker_slot_config (
    tenant_id,
    worker_id,
    slot_type,
    max_units,
    created_at,
    updated_at
)
SELECT
    @tenantId::uuid,
    @workerId::uuid,
    unnest(@slotTypes::text[]),
    unnest(@maxUnits::integer[]),
    CURRENT_TIMESTAMP,
    CURRENT_TIMESTAMP
-- NOTE: ON CONFLICT can be removed after the 0_76_d migration is run to remove insert triggers added in 0_76
ON CONFLICT (tenant_id, worker_id, slot_type) DO UPDATE SET
    max_units = EXCLUDED.max_units,
    updated_at = CURRENT_TIMESTAMP;

-- name: ListAvailableSlotsForWorkers :many
WITH worker_capacities AS (
    SELECT
        worker_id,
        max_units
    FROM
        v1_worker_slot_config
    WHERE
        tenant_id = @tenantId::uuid
        AND worker_id = ANY(@workerIds::uuid[])
        AND slot_type = @slotType::text
), worker_used_slots AS (
    SELECT
        worker_id,
        SUM(units) AS used_units
    FROM
        v1_task_runtime_slot
    WHERE
        tenant_id = @tenantId::uuid
        AND worker_id = ANY(@workerIds::uuid[])
        AND slot_type = @slotType::text
    GROUP BY
        worker_id
)
SELECT
    wc.worker_id AS "id",
    wc.max_units - COALESCE(wus.used_units, 0) AS "availableSlots"
FROM
    worker_capacities wc
LEFT JOIN
    worker_used_slots wus ON wc.worker_id = wus.worker_id;

-- name: ListAvailableSlotsForWorkersAndTypes :many
WITH worker_capacities AS (
    SELECT
        worker_id,
        slot_type,
        max_units
    FROM
        v1_worker_slot_config
    WHERE
        tenant_id = @tenantId::uuid
        AND worker_id = ANY(@workerIds::uuid[])
        AND slot_type = ANY(@slotTypes::text[])
), worker_used_slots AS (
    SELECT
        worker_id,
        slot_type,
        SUM(units) AS used_units
    FROM
        v1_task_runtime_slot
    WHERE
        tenant_id = @tenantId::uuid
        AND worker_id = ANY(@workerIds::uuid[])
        AND slot_type = ANY(@slotTypes::text[])
    GROUP BY
        worker_id,
        slot_type
)
SELECT
    wc.worker_id AS "id",
    wc.slot_type AS "slotType",
    wc.max_units - COALESCE(wus.used_units, 0) AS "availableSlots"
FROM
    worker_capacities wc
LEFT JOIN
    worker_used_slots wus ON wc.worker_id = wus.worker_id AND wc.slot_type = wus.slot_type;

-- name: ListWorkers :many
SELECT
    sqlc.embed(workers)
FROM
    "Worker" workers
WHERE
    workers."tenantId" = @tenantId
    AND (
        sqlc.narg('actionId')::text IS NULL OR
        workers."id" IN (
            SELECT "_ActionToWorker"."B"
            FROM "_ActionToWorker"
            INNER JOIN "Action" ON "Action"."id" = "_ActionToWorker"."A"
            WHERE "Action"."tenantId" = @tenantId AND "Action"."actionId" = sqlc.narg('actionId')::text
        )
    )
    AND (
        sqlc.narg('lastHeartbeatAfter')::timestamp IS NULL OR
        workers."lastHeartbeatAt" > sqlc.narg('lastHeartbeatAfter')::timestamp
    )
    AND (
        sqlc.narg('assignable')::boolean IS NULL OR
        (sqlc.narg('assignable')::boolean AND (
            SELECT COALESCE(SUM(cap.max_units), 0)
            FROM v1_worker_slot_config cap
            WHERE cap.tenant_id = workers."tenantId" AND cap.worker_id = workers."id"
        ) > (
            SELECT COALESCE(SUM(runtime.units), 0)
            FROM v1_task_runtime_slot runtime
            WHERE runtime.tenant_id = workers."tenantId" AND runtime.worker_id = workers."id"
        ))
    )
GROUP BY
    workers."id";

-- name: GetWorkerById :one
SELECT
    sqlc.embed(w)
FROM
    "Worker" w
WHERE
    w."id" = @id::uuid;

-- name: GetActiveWorkerById :one
SELECT
    sqlc.embed(w),
    ww."url" AS "webhookUrl",
    w."maxRuns" - (
        SELECT COUNT(*)
        FROM v1_task_runtime runtime
        WHERE
            runtime.tenant_id = w."tenantId" AND
            runtime.worker_id = w."id"
    ) AS "remainingSlots"
FROM
    "Worker" w
LEFT JOIN
    "WebhookWorker" ww ON w."webhookId" = ww."id"
WHERE
    w."id" = @id::uuid
    AND w."tenantId" = @tenantId::uuid
    AND w."dispatcherId" IS NOT NULL
    AND w."lastHeartbeatAt" > NOW() - INTERVAL '5 seconds'
    AND w."isActive" = true
    AND w."isPaused" = false;

-- name: ListSemaphoreSlotsWithStateForWorker :many
SELECT
    *
FROM
    v1_task_runtime runtime
JOIN
    v1_task ON runtime.task_id = v1_task.id AND runtime.task_inserted_at = v1_task.inserted_at
WHERE
    runtime.tenant_id = @tenantId::uuid
    AND runtime.worker_id = @workerId::uuid
LIMIT
    COALESCE(sqlc.narg('limit')::int, 100);

-- name: ListTotalActiveSlotsPerTenant :many
SELECT
    wc.tenant_id AS "tenantId",
    SUM(wc.max_units) AS "totalActiveSlots"
FROM v1_worker_slot_config wc
JOIN "Worker" w ON w."id" = wc.worker_id AND w."tenantId" = wc.tenant_id
WHERE
    w."dispatcherId" IS NOT NULL
    AND w."lastHeartbeatAt" > NOW() - INTERVAL '5 seconds'
    AND w."isActive" = true
    AND w."isPaused" = false
GROUP BY wc.tenant_id
;

-- name: ListActiveSlotsPerTenantAndSlotType :many
SELECT
    wc.tenant_id AS "tenantId",
    wc.slot_type AS "slotType",
    SUM(wc.max_units) AS "activeSlots"
FROM v1_worker_slot_config wc
JOIN "Worker" w ON w."id" = wc.worker_id AND w."tenantId" = wc.tenant_id
WHERE
    w."dispatcherId" IS NOT NULL
    AND w."lastHeartbeatAt" > NOW() - INTERVAL '5 seconds'
    AND w."isActive" = true
    AND w."isPaused" = false
GROUP BY wc.tenant_id, wc.slot_type
;

-- name: ListActiveSDKsPerTenant :many
SELECT
    "tenantId",
    COALESCE("language"::TEXT, 'unknown')::TEXT AS "language",
    COALESCE("languageVersion", 'unknown') AS "languageVersion",
    COALESCE("sdkVersion", 'unknown') AS "sdkVersion",
    COALESCE("os", 'unknown') AS "os",
    COUNT(*) AS "count"
FROM "Worker"
WHERE
    "dispatcherId" IS NOT NULL
    AND "lastHeartbeatAt" > NOW() - INTERVAL '5 seconds'
    AND "isActive" = true
    AND "isPaused" = false
GROUP BY "tenantId", "language", "languageVersion", "sdkVersion", "os"
;

-- name: ListActiveWorkersPerTenant :many
SELECT "tenantId", COUNT(*)
FROM "Worker"
WHERE
    "dispatcherId" IS NOT NULL
    AND "lastHeartbeatAt" > NOW() - INTERVAL '5 seconds'
    AND "isActive" = true
    AND "isPaused" = false
GROUP BY "tenantId"
;

-- name: GetWorkerActionsByWorkerId :many
SELECT
    w."id" AS "workerId",
    a."actionId" AS actionId
FROM "Worker" w
JOIN "_ActionToWorker" aw ON w.id = aw."B"
JOIN "Action" a ON aw."A" = a.id
WHERE
    a."tenantId" = @tenantId::UUID
    AND w.id = ANY(@workerIds::UUID[])
;

-- name: GetWorkerActionsByWorkerActionHash :many
SELECT DISTINCT ON (w."actionHash")
    w."actionHash" AS "actionHash",
    a."actionId" AS actionId
FROM "Worker" w
JOIN "_ActionToWorker" aw ON w.id = aw."B"
JOIN "Action" a ON aw."A" = a.id
WHERE
    a."tenantId" = @tenantId::UUID
    AND w."actionHash" = ANY(@actionHashes::BYTEA[])
ORDER BY w."actionHash", a."createdAt" DESC
;

-- name: GetWorkerWorkflowsByWorkerId :many
SELECT wf.*
FROM "Worker" w
JOIN "_ActionToWorker" aw ON w.id = aw."B"
JOIN "Action" a ON aw."A" = a.id
JOIN "Step" s ON s."actionId" = a."actionId" AND s."tenantId" = a."tenantId"
JOIN "Job" j ON j."id" = s."jobId" AND j."tenantId" = a."tenantId"
JOIN "WorkflowVersion" wv ON wv."id" = j."workflowVersionId"
JOIN "Workflow" wf ON wf."id" = wv."workflowId" AND wf."tenantId" = a."tenantId"
WHERE
    w."id" = @workerId::UUID
    AND w."tenantId" = @tenantId::UUID
;

-- name: GetWorkerForEngine :one
SELECT
    w."id" AS "id",
    w."tenantId" AS "tenantId",
    w."dispatcherId" AS "dispatcherId",
    w."lastHeartbeatAt" AS "lastHeartbeatAt",
    d."lastHeartbeatAt" AS "dispatcherLastHeartbeatAt",
    w."isActive" AS "isActive",
    w."lastListenerEstablished" AS "lastListenerEstablished"
FROM
    "Worker" w
LEFT JOIN
    "Dispatcher" d ON w."dispatcherId" = d."id"
WHERE
    w."tenantId" = @tenantId
    AND w."id" = @id;

-- name: ListWorkerLabels :many
SELECT
    "workerId",
    "id",
    "key",
    "intValue",
    "strValue",
    "createdAt",
    "updatedAt"
FROM "WorkerLabel" wl
WHERE wl."workerId" = ANY(@workerIds::uuid[]);

-- name: UpdateWorker :one
UPDATE
    "Worker"
SET
    "updatedAt" = CURRENT_TIMESTAMP,
    "dispatcherId" = coalesce(sqlc.narg('dispatcherId')::uuid, "dispatcherId"),
    "lastHeartbeatAt" = coalesce(sqlc.narg('lastHeartbeatAt')::timestamp, "lastHeartbeatAt"),
    "isActive" = coalesce(sqlc.narg('isActive')::boolean, "isActive"),
    "isPaused" = coalesce(sqlc.narg('isPaused')::boolean, "isPaused")
WHERE
    "id" = @id::uuid
RETURNING *;

-- name: LinkActionsToWorker :exec
INSERT INTO "_ActionToWorker" (
    "A",
    "B"
) SELECT
    unnest(@actionIds::uuid[]),
    @workerId::uuid
ON CONFLICT DO NOTHING;

-- name: UpdateWorkerHeartbeat :one
UPDATE
    "Worker"
SET
    "updatedAt" = CURRENT_TIMESTAMP,
    "lastHeartbeatAt" = sqlc.narg('lastHeartbeatAt')::timestamp
WHERE
    "id" = @id::uuid
RETURNING *;

-- name: DeleteWorker :one
DELETE FROM
  "Worker"
WHERE
  "id" = @id::uuid
RETURNING *;

-- name: UpdateWorkerActiveStatus :one
UPDATE "Worker"
SET
    "isActive" = @isActive::boolean,
    "lastListenerEstablished" = sqlc.narg('lastListenerEstablished')::timestamp
WHERE
    "id" = @id::uuid
    AND (
        "lastListenerEstablished" IS NULL
        OR "lastListenerEstablished" <= sqlc.narg('lastListenerEstablished')::timestamp
        )
RETURNING *;

-- name: UpsertWorkerLabel :one
INSERT INTO "WorkerLabel" (
    "createdAt",
    "updatedAt",
    "workerId",
    "key",
    "intValue",
    "strValue"
) VALUES (
    CURRENT_TIMESTAMP,
    CURRENT_TIMESTAMP,
    @workerId::uuid,
    @key::text,
    sqlc.narg('intValue')::int,
    sqlc.narg('strValue')::text
) ON CONFLICT ("workerId", "key") DO UPDATE
SET
    "updatedAt" = CURRENT_TIMESTAMP,
    "intValue" = sqlc.narg('intValue')::int,
    "strValue" = sqlc.narg('strValue')::text
RETURNING *;

-- name: CleanupOldWorkers :execresult
WITH old_workers AS (
    SELECT "id"
    FROM "Worker"
    WHERE "tenantId" = @tenantId::uuid
      AND "lastHeartbeatAt" < @lastHeartbeatBefore::timestamp
    LIMIT @batchSize::int
), deleted_worker_slot_configs AS (
    DELETE FROM v1_worker_slot_config
    WHERE worker_id IN (SELECT "id" FROM old_workers)
)
DELETE FROM "Worker"
WHERE "id" IN (SELECT "id" FROM old_workers);

-- name: ListDispatcherIdsForWorkers :many
SELECT
    "id" as "workerId",
    "dispatcherId"
FROM
    "Worker"
WHERE
    "tenantId" = @tenantId::uuid
    AND "id" = ANY(@workerIds::uuid[]);

-- name: UpsertService :one
INSERT INTO "Service" (
    "id",
    "createdAt",
    "updatedAt",
    "name",
    "tenantId"
)
VALUES (
    gen_random_uuid(),
    CURRENT_TIMESTAMP,
    CURRENT_TIMESTAMP,
    @name::text,
    @tenantId::uuid
)
ON CONFLICT ("tenantId", "name") DO UPDATE
SET
    "updatedAt" = CURRENT_TIMESTAMP
WHERE
    "Service"."tenantId" = @tenantId AND "Service"."name" = @name::text
RETURNING *;

-- name: CreateWorker :one
INSERT INTO "Worker" (
    "id",
    "createdAt",
    "updatedAt",
    "tenantId",
    "name",
    "dispatcherId",
    "type",
    "sdkVersion",
    "language",
    "languageVersion",
    "os",
    "runtimeExtra",
    "actionHash"
) VALUES (
    gen_random_uuid(),
    CURRENT_TIMESTAMP,
    CURRENT_TIMESTAMP,
    @tenantId::uuid,
    @name::text,
    @dispatcherId::uuid,
    sqlc.narg('type')::"WorkerType",
    sqlc.narg('sdkVersion')::text,
    sqlc.narg('language')::"WorkerSDKS",
    sqlc.narg('languageVersion')::text,
    sqlc.narg('os')::text,
    sqlc.narg('runtimeExtra')::text,
    @actionHash::bytea
) RETURNING *;

-- name: LinkServicesToWorker :exec
INSERT INTO "_ServiceToWorker" (
    "A",
    "B"
)
VALUES (
    unnest(@services::uuid[]),
    @workerId::uuid
)
ON CONFLICT DO NOTHING;

-- name: UpdateWorkerDurableTaskDispatcherId :one
UPDATE "Worker"
SET
    "durableTaskDispatcherId" = @dispatcherId::UUID,
    "updatedAt" = CURRENT_TIMESTAMP
WHERE
    "id" = @workerId::uuid
    AND "tenantId" = @tenantId::uuid
RETURNING *;

-- name: ListDurableTaskDispatcherIdsForTasks :many
WITH tasks AS (
    SELECT
        UNNEST(@taskIds::BIGINT[]) AS task_id,
        UNNEST(@taskInsertedAts::TIMESTAMPTZ[]) AS task_inserted_at
)

SELECT
    rt.*,
    w."durableTaskDispatcherId"
FROM v1_task_runtime rt
LEFT JOIN "Worker" w ON rt.worker_id = w.id
WHERE
    rt.tenant_id = @tenantId::uuid
    AND (rt.task_id, rt.task_inserted_at) IN (
        SELECT task_id, task_inserted_at
        FROM tasks
    )
;