# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: mark_running_tasks_lost.sql
import dataclasses
import uuid

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


MARK_RUNNING_TASKS_LOST = """-- name: mark_running_tasks_lost \\:exec
WITH lost_tasks AS (
    UPDATE hyrex_task_run
    SET status = 'LOST'\\:\\:task_run_status,
        finished = CURRENT_TIMESTAMP
    WHERE executor_id = :p1\\:\\:UUID
      AND status = 'RUNNING'\\:\\:task_run_status
    RETURNING id, durable_id, root_id, parent_id, workflow_run_id, workflow_dependencies, task_name, args, queue, max_retries, priority, timeout_seconds, status, attempt_number, scheduled_start, executor_id, queued, started, finished, last_heartbeat, idempotency_key, log_link
)
INSERT INTO hyrex_task_run (
    id,
    durable_id,
    root_id,
    parent_id,
    workflow_run_id,
    workflow_dependencies,
    task_name,
    args,
    queue,
    max_retries,
    priority,
    timeout_seconds,
    status,
    attempt_number,
    idempotency_key,
    queued
)
SELECT
    uuid7(),
    durable_id,
    root_id,
    parent_id,
    workflow_run_id,
    workflow_dependencies,
    task_name,
    args,
    queue,
    max_retries,
    priority,
    timeout_seconds,
    'QUEUED'\\:\\:task_run_status,
    attempt_number + 1,
    NULL, -- No idempotency key for retries
    CURRENT_TIMESTAMP
FROM lost_tasks
WHERE attempt_number < max_retries
"""


@dataclasses.dataclass()
class MarkRunningTasksLostParams:
    executor_id: uuid.UUID


class Querier:
    def __init__(self, conn: sqlalchemy.engine.Connection):
        self._conn = conn

    def mark_running_tasks_lost(self, arg: MarkRunningTasksLostParams) -> None:
        self._conn.execute(sqlalchemy.text(MARK_RUNNING_TASKS_LOST), {"p1": arg.executor_id})


class AsyncQuerier:
    def __init__(self, conn: sqlalchemy.ext.asyncio.AsyncConnection):
        self._conn = conn

    async def mark_running_tasks_lost(self, arg: MarkRunningTasksLostParams) -> None:
        await self._conn.execute(sqlalchemy.text(MARK_RUNNING_TASKS_LOST), {"p1": arg.executor_id})
