karldns.de/karldns/store/__init__.py

from __future__ import annotations

import threading
from collections.abc import Iterator
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) -> Iterator[Connection]:
        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"]