stable/freqtrade/persistence/pairlock.py

from datetime import datetime, timezone
from typing import Any, ClassVar, Dict, Optional

from sqlalchemy import ScalarResult, String, or_, select
from sqlalchemy.orm import Mapped, mapped_column

from freqtrade.constants import DATETIME_PRINT_FORMAT
from freqtrade.persistence.base import ModelBase, SessionType


class PairLock(ModelBase):
    """
    Pair Locks database model.
    """
    __tablename__ = 'pairlocks'
    session: ClassVar[SessionType]

    id: Mapped[int] = mapped_column(primary_key=True)

    pair: Mapped[str] = mapped_column(String(25), nullable=False, index=True)
    # lock direction - long, short or * (for both)
    side: Mapped[str] = mapped_column(String(25), nullable=False, default="*")
    reason: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
    # Time the pair was locked (start time)
    lock_time: Mapped[datetime] = mapped_column(nullable=False)
    # Time until the pair is locked (end time)
    lock_end_time: Mapped[datetime] = mapped_column(nullable=False, index=True)

    active: Mapped[bool] = mapped_column(nullable=False, default=True, index=True)

    def __repr__(self) -> str:
        lock_time = self.lock_time.strftime(DATETIME_PRINT_FORMAT)
        lock_end_time = self.lock_end_time.strftime(DATETIME_PRINT_FORMAT)
        return (
            f'PairLock(id={self.id}, pair={self.pair}, side={self.side}, lock_time={lock_time}, '
            f'lock_end_time={lock_end_time}, reason={self.reason}, active={self.active})')

    @staticmethod
    def query_pair_locks(
            pair: Optional[str], now: datetime, side: str = '*') -> ScalarResult['PairLock']:
        """
        Get all currently active locks for this pair
        :param pair: Pair to check for. Returns all current locks if pair is empty
        :param now: Datetime object (generated via datetime.now(timezone.utc)).
        """
        filters = [PairLock.lock_end_time > now,
                   # Only active locks
                   PairLock.active.is_(True), ]
        if pair:
            filters.append(PairLock.pair == pair)
        if side != '*':
            filters.append(or_(PairLock.side == side, PairLock.side == '*'))
        else:
            filters.append(PairLock.side == '*')

        return PairLock.session.scalars(select(PairLock).filter(*filters))

    @staticmethod
    def get_all_locks() -> ScalarResult['PairLock']:
        return PairLock.session.scalars(select(PairLock))

    def to_json(self) -> Dict[str, Any]:
        return {
            'id': self.id,
            'pair': self.pair,
            'lock_time': self.lock_time.strftime(DATETIME_PRINT_FORMAT),
            'lock_timestamp': int(self.lock_time.replace(tzinfo=timezone.utc).timestamp() * 1000),
            'lock_end_time': self.lock_end_time.strftime(DATETIME_PRINT_FORMAT),
            'lock_end_timestamp': int(self.lock_end_time.replace(tzinfo=timezone.utc
                                                                 ).timestamp() * 1000),
            'reason': self.reason,
            'side': self.side,
            'active': self.active,
        }