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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


GET_WORKFLOW_RUN_TASK_RUNS = """-- name: get_workflow_run_task_runs \\:many
WITH latest_attempts AS (
    SELECT 
        id as task_id,
        durable_id,
        task_name,
        status,
        workflow_dependencies,
        queued,
        ROW_NUMBER() OVER (PARTITION BY durable_id ORDER BY queued DESC) as rn
    FROM hyrex_task_run
    WHERE workflow_run_id = :p1
)
SELECT 
    task_id,
    durable_id,
    task_name,
    status,
    workflow_dependencies
FROM latest_attempts
WHERE rn = 1
ORDER BY queued
"""


@dataclasses.dataclass()
class GetWorkflowRunTaskRunsParams:
    workflow_run_id: Optional[uuid.UUID]


@dataclasses.dataclass()
class GetWorkflowRunTaskRunsRow:
    task_id: uuid.UUID
    durable_id: uuid.UUID
    task_name: str
    status: Any
    workflow_dependencies: Optional[List[uuid.UUID]]


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

    def get_workflow_run_task_runs(self, arg: GetWorkflowRunTaskRunsParams) -> Iterator[GetWorkflowRunTaskRunsRow]:
        result = self._conn.execute(sqlalchemy.text(GET_WORKFLOW_RUN_TASK_RUNS), {"p1": arg.workflow_run_id})
        for row in result:
            yield GetWorkflowRunTaskRunsRow(
                task_id=row[0],
                durable_id=row[1],
                task_name=row[2],
                status=row[3],
                workflow_dependencies=row[4],
            )


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

    async def get_workflow_run_task_runs(self, arg: GetWorkflowRunTaskRunsParams) -> AsyncIterator[GetWorkflowRunTaskRunsRow]:
        result = await self._conn.stream(sqlalchemy.text(GET_WORKFLOW_RUN_TASK_RUNS), {"p1": arg.workflow_run_id})
        async for row in result:
            yield GetWorkflowRunTaskRunsRow(
                task_id=row[0],
                durable_id=row[1],
                task_name=row[2],
                status=row[3],
                workflow_dependencies=row[4],
            )
