~liberachat/beryllia

ref: 9eb98b89a15f8a920fba58bf6a956e27451f76ef beryllia/beryllia/database/kline_tag.py -rw-r--r-- 2.3 KiB
9eb98b89 — jesopo fix klinedel regex 3 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from datetime import datetime
from typing import Collection, Tuple

from .common import Table
from ..normalise import SearchType
from ..util import glob_to_sql, lex_glob_pattern


class KLineTagTable(Table):
    async def add(self, kline_id: int, tag: str, source: str, oper: str):
        query = """
            INSERT INTO kline_tag
                (kline_id, tag, search_tag, source, oper, ts)
            VALUES ($1, $2, $3, $4, $5, NOW()::TIMESTAMP)
        """
        async with self.pool.acquire() as conn:
            await conn.execute(
                query,
                kline_id,
                tag,
                str(self.to_search(tag, SearchType.TAG)),
                source,
                oper,
            )

    async def remove(self, kline_id: int, tag: str):
        query = """
            DELETE FROM kline_tag
            WHERE kline_id = $1
            AND search_tag = $2
        """
        async with self.pool.acquire() as conn:
            await conn.execute(
                query, kline_id, str(self.to_search(tag, SearchType.TAG))
            )

    async def exists(self, kline_id: int, tag: str) -> bool:
        query = """
            SELECT 1 FROM kline_tag
            WHERE kline_id = $1
            AND search_tag = $2
        """

        async with self.pool.acquire() as conn:
            return bool(
                await conn.fetchval(
                    query, kline_id, str(self.to_search(tag, SearchType.TAG))
                )
            )

    async def find_klines(
        self, tag: str, count: int
    ) -> Collection[Tuple[int, datetime]]:
        query = """
            SELECT DISTINCT(kline.id), kline.ts
                FROM kline_tag
            INNER JOIN kline
                ON kline_tag.kline_id = kline.id
            WHERE kline_tag.search_tag LIKE $1
            LIMIT $2
        """
        pattern = glob_to_sql(lex_glob_pattern(tag))
        param = str(self.to_search(pattern, SearchType.TAG))
        async with self.pool.acquire() as conn:
            return await conn.fetch(query, param, count)

    async def find_tags(self, kline_id: int) -> Collection[str]:
        query = """
            SELECT tag
            FROM kline_tag
            WHERE kline_id = $1
        """

        async with self.pool.acquire() as conn:
            rows = await conn.execute(query, kline_id)
        return {r[0] for r in rows}