Prepare KarlDNS production deployment

karldns/store/__init__.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import threading
+from contextlib import contextmanager
+from pathlib import Path
+from sqlite3 import Connection
+from time import monotonic
+
+from .accounts import AccountsMixin
+from .metrics import MetricsMixin
+from .rate_limit import RateLimitCleanupMixin
+from .records import RecordsEventsMixin
+from .scheduler import SchedulerMixin
+from .schema import connect_database, initialize_schema
+
+
+class DdnsStore(AccountsMixin, MetricsMixin, RecordsEventsMixin, RateLimitCleanupMixin, SchedulerMixin):
+    def __init__(self, database_path: Path) -> None:
+        self.database_path = database_path
+        self._connections = threading.local()
+        self._connection_lock = threading.Lock()
+        self._tracked_connections: list[Connection] = []
+        self._closed_connection_ids: set[int] = set()
+        self._rate_limit_prune_lock = threading.Lock()
+        self._rate_limit_lock = threading.Lock()
+        self._rate_limit_buckets: dict[str, dict[str, float | int]] = {}
+        self._rate_limit_dirty_counts: dict[tuple[str, int], int] = {}
+        self._last_rate_limit_flush = monotonic()
+        self._last_rate_limit_prune_window = -1
+        self.database_path.parent.mkdir(parents=True, exist_ok=True)
+        self._init_db()
+
+    @contextmanager
+    def connect(self):
+        conn = self._thread_connection()
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+
+    def _init_db(self) -> None:
+        with self.connect() as conn:
+            initialize_schema(conn)
+
+    def optimize(self) -> None:
+        with self.connect() as conn:
+            conn.execute("PRAGMA optimize")
+
+    def checkpoint_wal(self) -> None:
+        if str(self.database_path) == ":memory:":
+            return
+        with self.connect() as conn:
+            conn.execute("PRAGMA wal_checkpoint(PASSIVE)")
+
+    def _thread_connection(self) -> Connection:
+        conn = getattr(self._connections, "conn", None)
+        if conn is None or id(conn) in self._closed_connection_ids:
+            conn = connect_database(self.database_path)
+            self._connections.conn = conn
+            with self._connection_lock:
+                self._tracked_connections.append(conn)
+                self._closed_connection_ids.discard(id(conn))
+        return conn
+
+    def close_connections(self) -> None:
+        with self._connection_lock:
+            connections = list(self._tracked_connections)
+            self._tracked_connections.clear()
+            self._closed_connection_ids.update(id(conn) for conn in connections)
+        for conn in connections:
+            conn.close()
+        if getattr(self._connections, "conn", None) in connections:
+            self._connections.conn = None
+
+    close = close_connections
+
+
+__all__ = ["DdnsStore"]