# Code generated by sqlc. DO NOT EDIT.
# versions:
#   sqlc v1.29.0
# source: acquire_scheduler_lock.sql
import dataclasses
import datetime
from typing import Optional

import sqlalchemy
import sqlalchemy.ext.asyncio

from . import models


ACQUIRE_SCHEDULER_LOCK = """-- name: acquire_scheduler_lock \\:one
INSERT INTO hyrex_scheduler_lock (
    lockid, worker_name, acquired_at, heartbeat_at, release_at, is_active
)
VALUES (
    1,
    :p1,
    now(),
    now(),
    now() + CAST(:p2 AS interval),
    true
)
ON CONFLICT (lockid)
  DO UPDATE 
     SET worker_name  = EXCLUDED.worker_name,
         acquired_at  = EXCLUDED.acquired_at,
         heartbeat_at = EXCLUDED.heartbeat_at,
         release_at   = EXCLUDED.release_at,
         is_active    = EXCLUDED.is_active
   WHERE (
       -- Only update if the lock is not truly active,
       -- i.e. is already inactive or expired\\:
       hyrex_scheduler_lock.is_active = false
       OR hyrex_scheduler_lock.release_at <= now()
   )
RETURNING lockid
"""


@dataclasses.dataclass()
class AcquireSchedulerLockParams:
    worker_name: str
    duration: datetime.timedelta


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

    def acquire_scheduler_lock(self, arg: AcquireSchedulerLockParams) -> Optional[int]:
        row = self._conn.execute(sqlalchemy.text(ACQUIRE_SCHEDULER_LOCK), {"p1": arg.worker_name, "p2": arg.duration}).first()
        if row is None:
            return None
        return row[0]


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

    async def acquire_scheduler_lock(self, arg: AcquireSchedulerLockParams) -> Optional[int]:
        row = (await self._conn.execute(sqlalchemy.text(ACQUIRE_SCHEDULER_LOCK), {"p1": arg.worker_name, "p2": arg.duration})).first()
        if row is None:
            return None
        return row[0]
