Archived
82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
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"]
|