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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


SET_ORPHANED_TASK_EXECUTION_TO_LOST_AND_RETRY = """-- name: set_orphaned_task_execution_to_lost_and_retry \\:exec
WITH lost_tasks AS (
    UPDATE hyrex_task_run
        SET status = 'LOST'\\:\\:task_run_status
        WHERE status = 'RUNNING'\\:\\:task_run_status
            AND (
                  executor_id IS NULL
                      OR NOT EXISTS (
                      SELECT 1
                      FROM hyrex_executor
                      WHERE hyrex_executor.id = hyrex_task_run.executor_id
                        AND hyrex_executor.status = 'RUNNING'\\:\\:executor_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
    gen_random_uuid(),
    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,
    idempotency_key,
    NOW()
FROM lost_tasks
WHERE attempt_number < max_retries
"""


@dataclasses.dataclass()
class SetOrphanedTaskExecutionToLostAndRetryParams:

    pass

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

    def set_orphaned_task_execution_to_lost_and_retry(self, arg: SetOrphanedTaskExecutionToLostAndRetryParams) -> None:
        self._conn.execute(sqlalchemy.text(SET_ORPHANED_TASK_EXECUTION_TO_LOST_AND_RETRY))


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

    async def set_orphaned_task_execution_to_lost_and_retry(self, arg: SetOrphanedTaskExecutionToLostAndRetryParams) -> None:
        await self._conn.execute(sqlalchemy.text(SET_ORPHANED_TASK_EXECUTION_TO_LOST_AND_RETRY))
