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