ligo.at
decentralized and customizable links page on top of atproto
ligo.at
atproto
link-in-bio
python
uv
nauta.one
1import sqlite3
2from logging import Logger
3from sqlite3 import Connection
4from typing import Generic, Literal, cast, override
5
6from flask import Flask, g
7
8from src.atproto.kv import KV as BaseKV
9from src.atproto.kv import K, V
10
11
12class KV(BaseKV, Generic[K, V]):
13 db: Connection
14 logger: Logger
15 prefix: str
16
17 def __init__(self, app: Connection | Flask, logger: Logger, prefix: str):
18 self.db = app if isinstance(app, Connection) else get_db(app, name="keyval")
19 self.logger = logger
20 self.prefix = prefix
21
22 @override
23 def get(self, key: K) -> V | None:
24 cursor = self.db.cursor()
25 row: dict[str, str] | None = cursor.execute(
26 "select value from keyval where prefix = ? and key = ?",
27 (self.prefix, key),
28 ).fetchone()
29 if row is not None:
30 self.logger.debug(f"returning cached {self.prefix}({key})")
31 return cast(V, row["value"])
32 return None
33
34 @override
35 def set(self, key: K, value: V):
36 self.logger.debug(f"caching {self.prefix}({key}): {value}")
37 cursor = self.db.cursor()
38 _ = cursor.execute(
39 "insert or replace into keyval (prefix, key, value) values (?, ?, ?)",
40 (self.prefix, key, value),
41 )
42 self.db.commit()
43
44
45type DatabaseName = Literal["config"] | Literal["keyval"]
46
47
48def get_db(app: Flask, name: DatabaseName) -> sqlite3.Connection:
49 global_key = f"{name}_db"
50 db: sqlite3.Connection | None = g.get(global_key, None)
51 if db is None:
52 db_path: str = app.config.get(f"{name.upper()}_DB_URL", f"{name}.db")
53 db = sqlite3.connect(db_path, check_same_thread=False)
54 setattr(g, global_key, db)
55 # return rows as dict-like objects
56 db.row_factory = sqlite3.Row
57 return db
58
59
60def close_db_connection(_exception: BaseException | None):
61 for name in ["keyval", "config"]:
62 db: sqlite3.Connection | None = g.pop(f"{name}_db", None)
63 if db is not None:
64 db.close()
65
66
67def init_db(app: Flask, name: DatabaseName) -> None:
68 with app.app_context():
69 db = get_db(app, name)
70 with app.open_resource(f"{name}.sql", mode="r") as schema:
71 _ = db.cursor().executescript(schema.read())
72 db.commit()