query.py

"""Query interface."""
from __future__ import annotations

import abc

import sqlalchemy.orm.query

from crawler.constants import RunStatus


class BaseQuery(metaclass=abc.ABCMeta):
    """Base query builder."""

    @abc.abstractmethod
    def query(self):
        """Return encapsulated query."""

    @abc.abstractmethod
    def limit(self, limit: int) -> BaseQuery:
        """Limit for entities."""

    @abc.abstractmethod
    def offset(self, offset: int) -> BaseQuery:
        """Offset for entities."""


class BaseDBQuery(metaclass=abc.ABCMeta):
    """Add db related features."""

    @abc.abstractmethod
    def lock(self) -> BaseRunQuery:
        """Lock runs in case if it's a db query."""


class BaseRunQuery(BaseQuery):
    """Base query for Run entity."""

    @abc.abstractmethod
    def status(self, status: RunStatus) -> BaseRunQuery:
        """Filter runs by status."""


class AlchemyRunQuery(BaseDBQuery, BaseQuery):
    """Run query that works with SQLAlchemy session."""

    def __init__(self, query: sqlalchemy.orm.Query):
        """Primary constructor.

        Args:
            query: instance of SQLAlchemy query.
        """
        self._query = query

    def query(self) -> sqlalchemy.orm.Query:
        """Return SQL Alchemy query."""
        return self._query

    def limit(self, limit: int):
        """Limit for runs."""
        return AlchemyRunQuery(self._query.limit(limit))

    def offset(self, offset: int):
        """Offset for runs."""
        return AlchemyRunQuery(self._query.offset(offset))

    def status(self, status: RunStatus):
        """Filter runs by status."""
        return AlchemyRunQuery(self._query.filter_by(status=status))

    def lock(self):
        """Lock rows in the db."""
        return self