chia-blockchain/chia/wallet/notification_store.py

from __future__ import annotations

import dataclasses
import logging
import sqlite3
from typing import List, Optional, Tuple

from chia.types.blockchain_format.sized_bytes import bytes32
from chia.util.db_wrapper import DBWrapper2
from chia.util.ints import uint32, uint64


@dataclasses.dataclass(frozen=True)
class Notification:
    coin_id: bytes32
    message: bytes
    amount: uint64
    height: uint32


class NotificationStore:
    """
    NotificationStore stores trading history.
    """

    cache_size: uint32
    db_wrapper: DBWrapper2
    log: logging.Logger

    @classmethod
    async def create(
        cls, db_wrapper: DBWrapper2, cache_size: uint32 = uint32(600000), name: Optional[str] = None
    ) -> "NotificationStore":
        self = cls()

        if name:
            self.log = logging.getLogger(name)
        else:
            self.log = logging.getLogger(__name__)

        self.cache_size = cache_size
        self.db_wrapper = db_wrapper

        async with self.db_wrapper.writer_maybe_transaction() as conn:
            await conn.execute(
                "CREATE TABLE IF NOT EXISTS notifications(coin_id blob PRIMARY KEY, msg blob, amount blob)"
            )

            await conn.execute("CREATE TABLE IF NOT EXISTS all_notification_ids(coin_id blob PRIMARY KEY)")

            try:
                await conn.execute("ALTER TABLE notifications ADD COLUMN height bigint DEFAULT 0")
            except sqlite3.OperationalError as e:
                if "duplicate column" in e.args[0]:
                    pass  # ignore what is likely Duplicate column error
                else:
                    raise e

            # This used to be an accidentally created redundant index on coin_id which is already a primary key
            # We can remove this at some point in the future when it's unlikely this index still exists
            await conn.execute("DROP INDEX IF EXISTS coin_id_index")

        return self

    async def add_notification(self, notification: Notification) -> None:
        """
        Store Notification into DB
        """
        async with self.db_wrapper.writer_maybe_transaction() as conn:
            cursor = await conn.execute(
                "INSERT OR REPLACE INTO notifications (coin_id, msg, amount, height) VALUES(?, ?, ?, ?)",
                (
                    notification.coin_id,
                    notification.message,
                    bytes(notification.amount),
                    notification.height,
                ),
            )
            cursor = await conn.execute(
                "INSERT OR REPLACE INTO all_notification_ids (coin_id) VALUES(?)",
                (notification.coin_id,),
            )
            await cursor.close()

    async def get_notifications(self, coin_ids: List[bytes32]) -> List[Notification]:
        """
        Checks DB for Notification with id: id and returns it.
        """
        coin_ids_str_list = "("
        for _ in coin_ids:
            coin_ids_str_list += "?"
            coin_ids_str_list += ","
        coin_ids_str_list = coin_ids_str_list[:-1] if len(coin_ids_str_list) > 1 else "("
        coin_ids_str_list += ")"

        async with self.db_wrapper.reader_no_transaction() as conn:
            rows = await conn.execute_fetchall(
                f"SELECT * from notifications WHERE coin_id IN {coin_ids_str_list} ORDER BY amount DESC", coin_ids
            )

        return [
            Notification(
                bytes32(row[0]),
                bytes(row[1]),
                uint64.from_bytes(row[2]),
                uint32(row[3]),
            )
            for row in rows
        ]

    async def get_all_notifications(
        self, pagination: Optional[Tuple[Optional[int], Optional[int]]] = None
    ) -> List[Notification]:
        """
        Checks DB for Notification with id: id and returns it.
        """
        if pagination is not None:
            if pagination[1] is not None and pagination[0] is not None:
                pagination_str = f" LIMIT {pagination[0]}, {pagination[1] - pagination[0]}"
            elif pagination[1] is None and pagination[0] is not None:
                pagination_str = f" LIMIT {pagination[0]}, (SELECT COUNT(*) from notifications)"
            elif pagination[1] is not None and pagination[0] is None:
                pagination_str = f" LIMIT {pagination[1]}"
            else:
                pagination_str = ""
        else:
            pagination_str = ""

        async with self.db_wrapper.reader_no_transaction() as conn:
            rows = await conn.execute_fetchall(f"SELECT * from notifications ORDER BY amount DESC{pagination_str}")

        return [
            Notification(
                bytes32(row[0]),
                bytes(row[1]),
                uint64.from_bytes(row[2]),
                uint32(row[3]),
            )
            for row in rows
        ]

    async def delete_notifications(self, coin_ids: List[bytes32]) -> None:
        coin_ids_str_list = "("
        for _ in coin_ids:
            coin_ids_str_list += "?"
            coin_ids_str_list += ","
        coin_ids_str_list = coin_ids_str_list[:-1] if len(coin_ids_str_list) > 1 else "("
        coin_ids_str_list += ")"

        async with self.db_wrapper.writer_maybe_transaction() as conn:
            # Delete from storage
            cursor = await conn.execute(f"DELETE FROM notifications WHERE coin_id IN {coin_ids_str_list}", coin_ids)
            await cursor.close()

    async def delete_all_notifications(self) -> None:
        async with self.db_wrapper.writer_maybe_transaction() as conn:
            # Delete from storage
            cursor = await conn.execute("DELETE FROM notifications")
            await cursor.close()

    async def notification_exists(self, id: bytes32) -> bool:
        async with self.db_wrapper.reader_no_transaction() as conn:
            async with conn.execute(
                "SELECT EXISTS (SELECT 1 from all_notification_ids WHERE coin_id=?)", (id,)
            ) as cursor:
                row = await cursor.fetchone()
                assert row is not None
                exists: bool = row[0] > 0
                return exists