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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


MARK_LOST_EXECUTORS = """-- name: mark_lost_executors \\:many
UPDATE hyrex_executor
SET status = 'LOST',
    stopped = last_heartbeat
WHERE status = 'RUNNING'
AND last_heartbeat < (CURRENT_TIMESTAMP - :p1\\:\\:INTERVAL)
AND stopped IS NULL
RETURNING id
"""


@dataclasses.dataclass()
class MarkLostExecutorsParams:
    timeout: datetime.timedelta


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

    def mark_lost_executors(self, arg: MarkLostExecutorsParams) -> Iterator[uuid.UUID]:
        result = self._conn.execute(sqlalchemy.text(MARK_LOST_EXECUTORS), {"p1": arg.timeout})
        for row in result:
            yield row[0]


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

    async def mark_lost_executors(self, arg: MarkLostExecutorsParams) -> AsyncIterator[uuid.UUID]:
        result = await self._conn.stream(sqlalchemy.text(MARK_LOST_EXECUTORS), {"p1": arg.timeout})
        async for row in result:
            yield row[0]
