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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


GET_TASK_RUNS_BY_STATUS_PAGINATED = """-- name: get_task_runs_by_status_paginated \\:many
SELECT 
    ROW_NUMBER() OVER (ORDER BY queued DESC) as row_number,
    id,
    durable_id,
    root_id,
    parent_id,
    task_name,
    args,
    queue,
    max_retries,
    priority,
    status,
    attempt_number,
    scheduled_start,
    executor_id,
    queued,
    started,
    finished,
    log_link
FROM hyrex_task_run
WHERE status = :p1
ORDER BY queued DESC
LIMIT :p2 OFFSET :p3
"""


@dataclasses.dataclass()
class GetTaskRunsByStatusPaginatedParams:
    status: Any
    limit: int
    offset: int


@dataclasses.dataclass()
class GetTaskRunsByStatusPaginatedRow:
    row_number: int
    id: uuid.UUID
    durable_id: uuid.UUID
    root_id: uuid.UUID
    parent_id: Optional[uuid.UUID]
    task_name: str
    args: Any
    queue: str
    max_retries: int
    priority: int
    status: Any
    attempt_number: int
    scheduled_start: Optional[datetime.datetime]
    executor_id: Optional[uuid.UUID]
    queued: Optional[datetime.datetime]
    started: Optional[datetime.datetime]
    finished: Optional[datetime.datetime]
    log_link: Optional[str]


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

    def get_task_runs_by_status_paginated(self, arg: GetTaskRunsByStatusPaginatedParams) -> Iterator[GetTaskRunsByStatusPaginatedRow]:
        result = self._conn.execute(sqlalchemy.text(GET_TASK_RUNS_BY_STATUS_PAGINATED), {"p1": arg.status, "p2": arg.limit, "p3": arg.offset})
        for row in result:
            yield GetTaskRunsByStatusPaginatedRow(
                row_number=row[0],
                id=row[1],
                durable_id=row[2],
                root_id=row[3],
                parent_id=row[4],
                task_name=row[5],
                args=row[6],
                queue=row[7],
                max_retries=row[8],
                priority=row[9],
                status=row[10],
                attempt_number=row[11],
                scheduled_start=row[12],
                executor_id=row[13],
                queued=row[14],
                started=row[15],
                finished=row[16],
                log_link=row[17],
            )


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

    async def get_task_runs_by_status_paginated(self, arg: GetTaskRunsByStatusPaginatedParams) -> AsyncIterator[GetTaskRunsByStatusPaginatedRow]:
        result = await self._conn.stream(sqlalchemy.text(GET_TASK_RUNS_BY_STATUS_PAGINATED), {"p1": arg.status, "p2": arg.limit, "p3": arg.offset})
        async for row in result:
            yield GetTaskRunsByStatusPaginatedRow(
                row_number=row[0],
                id=row[1],
                durable_id=row[2],
                root_id=row[3],
                parent_id=row[4],
                task_name=row[5],
                args=row[6],
                queue=row[7],
                max_retries=row[8],
                priority=row[9],
                status=row[10],
                attempt_number=row[11],
                scheduled_start=row[12],
                executor_id=row[13],
                queued=row[14],
                started=row[15],
                finished=row[16],
                log_link=row[17],
            )
