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

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


FETCH_ACTIVE_QUEUE_NAMES = """-- name: fetch_active_queue_names \\:many
WITH distinct_queues AS (SELECT DISTINCT queue
                         FROM hyrex_task_run
                         WHERE status = 'QUEUED'
                           AND queue SIMILAR TO :p1),
     queue_count AS (SELECT COUNT(*) AS cnt
                     FROM distinct_queues)
SELECT queue
FROM (
         -- If count <= 100000, just select all queues
         SELECT dq.queue
         FROM distinct_queues dq,
              queue_count qc
         WHERE qc.cnt <= 100000

         UNION ALL

         -- If count > 100000, select a random subset
         SELECT queue
         FROM (SELECT dq.queue,
                      row_number() OVER (ORDER BY random()) AS rn
               FROM distinct_queues dq,
                    queue_count qc
               WHERE qc.cnt > 100000) sub
         WHERE rn <= 100000) final_result
"""


@dataclasses.dataclass()
class FetchActiveQueueNamesParams:
    queue_pattern: str


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

    def fetch_active_queue_names(self, arg: FetchActiveQueueNamesParams) -> Iterator[str]:
        result = self._conn.execute(sqlalchemy.text(FETCH_ACTIVE_QUEUE_NAMES), {"p1": arg.queue_pattern})
        for row in result:
            yield row[0]


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

    async def fetch_active_queue_names(self, arg: FetchActiveQueueNamesParams) -> AsyncIterator[str]:
        result = await self._conn.stream(sqlalchemy.text(FETCH_ACTIVE_QUEUE_NAMES), {"p1": arg.queue_pattern})
        async for row in result:
            yield row[0]
