Base models for RedisOM with expire time

extended_base_redis_model.py

__author__ = 'Khiem Doan'
__github__ = 'https://github.com/khiemdoan'
__email__ = 'doankhiem.crazy@gmail.com'

from datetime import datetime, timedelta
from typing import Self

from aredis_om import EmbeddedJsonModel as RedisEmbeddedJsonModel
from aredis_om import HashModel as RedisHashModel
from aredis_om import JsonModel as RedisJsonModel
from redis.asyncio.client import Pipeline, Redis


class ExpirableMixin:
    class Meta:
        database: Redis
        expire: int | timedelta | None  # total seconds

    async def expire(self, seconds: int|timedelta) -> None:
        if not hasattr(self.Meta, 'database'):
            return
        if not isinstance(self.Meta.database, Redis):
            return
        if isinstance(seconds, timedelta):
            seconds = int(seconds.total_seconds())
        await self.Meta.database.expire(self.key(), seconds)

    async def expire_at(self, time: datetime) -> None:
        now = datetime.now()
        delta = time - now
        await self.expire(delta)


class HashModel(RedisHashModel, ExpirableMixin):
    class Meta:
        database: Redis
        expire: int | timedelta | None  # total seconds

    async def save(self, pipeline: Pipeline|None = None) -> Self:
        await super().save(pipeline=pipeline)
        if hasattr(self.Meta, 'expire'):
            if isinstance(self.Meta.expire, (int, timedelta)):
                await self.expire(self.Meta.expire)
        return self


class JsonModel(RedisJsonModel, ExpirableMixin):
    class Meta:
        database: Redis
        expire: int | timedelta | None  # total seconds

    async def save(self, pipeline: Pipeline|None = None) -> Self:
        await super().save(pipeline=pipeline)
        if hasattr(self.Meta, 'expire'):
            if isinstance(self.Meta.expire, (int, timedelta)):
                await self.expire(self.Meta.expire)


class EmbeddedJsonModel(RedisEmbeddedJsonModel, ExpirableMixin):
    class Meta:
        database: Redis
        expire: int | timedelta | None  # total seconds

    async def save(self, pipeline: Pipeline|None = None) -> Self:
        await super().save(pipeline=pipeline)
        if hasattr(self.Meta, 'expire'):
            if isinstance(self.Meta.expire, (int, timedelta)):
                await self.expire(self.Meta.expire)