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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


CLAIM_QUEUED_HYPE_CRON_JOB_RUNS = """-- name: claim_queued_hype_cron_job_runs \\:many
  UPDATE hype_cron_job_run_details
  SET status = 'PROCESSING',
      start_time = NOW()
  WHERE runid IN (
      SELECT runid
      FROM hype_cron_job_run_details
      WHERE status = 'QUEUED'
        AND schedule_time <= NOW()
      ORDER BY schedule_time ASC
      LIMIT :p1
      FOR UPDATE SKIP LOCKED
  )
  RETURNING jobid, runid, command_type\\:\\:text, command_params, status\\:\\:text, schedule_time, start_time, end_time
"""


@dataclasses.dataclass()
class ClaimQueuedHypeCronJobRunsParams:
    limit: int


@dataclasses.dataclass()
class ClaimQueuedHypeCronJobRunsRow:
    jobid: int
    runid: int
    command_type: str
    command_params: Any
    status: str
    schedule_time: datetime.datetime
    start_time: Optional[datetime.datetime]
    end_time: Optional[datetime.datetime]


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

    def claim_queued_hype_cron_job_runs(self, arg: ClaimQueuedHypeCronJobRunsParams) -> Iterator[ClaimQueuedHypeCronJobRunsRow]:
        result = self._conn.execute(sqlalchemy.text(CLAIM_QUEUED_HYPE_CRON_JOB_RUNS), {"p1": arg.limit})
        for row in result:
            yield ClaimQueuedHypeCronJobRunsRow(
                jobid=row[0],
                runid=row[1],
                command_type=row[2],
                command_params=row[3],
                status=row[4],
                schedule_time=row[5],
                start_time=row[6],
                end_time=row[7],
            )


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

    async def claim_queued_hype_cron_job_runs(self, arg: ClaimQueuedHypeCronJobRunsParams) -> AsyncIterator[ClaimQueuedHypeCronJobRunsRow]:
        result = await self._conn.stream(sqlalchemy.text(CLAIM_QUEUED_HYPE_CRON_JOB_RUNS), {"p1": arg.limit})
        async for row in result:
            yield ClaimQueuedHypeCronJobRunsRow(
                jobid=row[0],
                runid=row[1],
                command_type=row[2],
                command_params=row[3],
                status=row[4],
                schedule_time=row[5],
                start_time=row[6],
                end_time=row[7],
            )
