# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: fetch_task.sql
import dataclasses
import datetime
from typing import Any, List, Optional
import uuid

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


FETCH_TASK = """-- name: fetch_task \\:one
WITH next_task AS (SELECT id
                   FROM hyrex_task_run
                   WHERE hyrex_task_run.queue = :p2
                     AND hyrex_task_run.status = 'QUEUED'\\:\\:task_run_status
                     AND hyrex_task_run.task_name = ANY(:p3\\:\\:VARCHAR[])
                   ORDER BY priority ASC, queued
                       FOR UPDATE SKIP LOCKED
                   LIMIT 1)
UPDATE hyrex_task_run AS ht
SET status      = 'RUNNING'\\:\\:task_run_status,
    started     = CURRENT_TIMESTAMP,
    executor_id = :p1
FROM next_task
WHERE ht.id = next_task.id
RETURNING ht.id
    , ht.durable_id
    , ht.root_id
    , ht.parent_id
    , ht.workflow_run_id
    , ht.task_name
    , ht.args
    , ht.queue
    , ht.attempt_number
    , ht.max_retries
    , ht.priority
    , ht.timeout_seconds
    , ht.scheduled_start
    , ht.queued
    , ht.started
"""


@dataclasses.dataclass()
class FetchTaskParams:
    executor_id: Optional[uuid.UUID]
    queue: str
    task_names: List[str]


@dataclasses.dataclass()
class FetchTaskRow:
    id: uuid.UUID
    durable_id: uuid.UUID
    root_id: uuid.UUID
    parent_id: Optional[uuid.UUID]
    workflow_run_id: Optional[uuid.UUID]
    task_name: str
    args: Any
    queue: str
    attempt_number: int
    max_retries: int
    priority: int
    timeout_seconds: Optional[int]
    scheduled_start: Optional[datetime.datetime]
    queued: Optional[datetime.datetime]
    started: Optional[datetime.datetime]


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

    def fetch_task(self, arg: FetchTaskParams) -> Optional[FetchTaskRow]:
        row = self._conn.execute(sqlalchemy.text(FETCH_TASK), {"p1": arg.executor_id, "p2": arg.queue, "p3": arg.task_names}).first()
        if row is None:
            return None
        return FetchTaskRow(
            id=row[0],
            durable_id=row[1],
            root_id=row[2],
            parent_id=row[3],
            workflow_run_id=row[4],
            task_name=row[5],
            args=row[6],
            queue=row[7],
            attempt_number=row[8],
            max_retries=row[9],
            priority=row[10],
            timeout_seconds=row[11],
            scheduled_start=row[12],
            queued=row[13],
            started=row[14],
        )


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

    async def fetch_task(self, arg: FetchTaskParams) -> Optional[FetchTaskRow]:
        row = (await self._conn.execute(sqlalchemy.text(FETCH_TASK), {"p1": arg.executor_id, "p2": arg.queue, "p3": arg.task_names})).first()
        if row is None:
            return None
        return FetchTaskRow(
            id=row[0],
            durable_id=row[1],
            root_id=row[2],
            parent_id=row[3],
            workflow_run_id=row[4],
            task_name=row[5],
            args=row[6],
            queue=row[7],
            attempt_number=row[8],
            max_retries=row[9],
            priority=row[10],
            timeout_seconds=row[11],
            scheduled_start=row[12],
            queued=row[13],
            started=row[14],
        )
