Add formatting and type checks

karldns/store/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import threading
+from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
 from sqlite3 import Connection
@@ -31,7 +32,7 @@ class DdnsStore(AccountsMixin, MetricsMixin, RecordsEventsMixin, RateLimitCleanu
         self._init_db()
 
     @contextmanager
-    def connect(self):
+    def connect(self) -> Iterator[Connection]:
         conn = self._thread_connection()
         try:
             yield conn

karldns/store/accounts.py

@@ -3,10 +3,11 @@ from __future__ import annotations
 from datetime import UTC, datetime
 
 from ..security import hash_lookup_token, hash_secret, new_bearer_secret, verify_secret
-from .types import AccountRow
+from .base import StoreConnectionMixin
+from .types import AccountRow, account_row
 
 
-class AccountsMixin:
+class AccountsMixin(StoreConnectionMixin):
     def create_account(self, hostname: str, username: str | None = None, *, accepted: bool = True) -> AccountRow:
         now = datetime.now(UTC).isoformat()
         update_slug = new_bearer_secret()
@@ -41,7 +42,7 @@ class AccountsMixin:
             """,
                 (limit, offset),
             ).fetchall()
-        return [dict(row) for row in rows]
+        return [account_row(row) for row in rows]
 
     def count_accounts(self) -> int:
         with self.connect() as conn:
@@ -60,7 +61,7 @@ class AccountsMixin:
                 """,
                 (hostname,),
             ).fetchone()
-        return dict(account) if account else None
+        return account_row(account) if account else None
 
     def get_account_by_slug(self, update_slug: str) -> AccountRow | None:
         with self.connect() as conn:
@@ -71,7 +72,7 @@ class AccountsMixin:
                 """,
                 (update_slug,),
             ).fetchone()
-        return dict(account) if account else None
+        return account_row(account) if account else None
 
     def get_account_by_management_slug(self, management_slug: str) -> AccountRow | None:
         with self.connect() as conn:
@@ -85,7 +86,7 @@ class AccountsMixin:
                 """,
                 (hash_lookup_token(management_slug),),
             ).fetchone()
-        return dict(account) if account else None
+        return account_row(account) if account else None
 
     def activate_account(self, hostname: str) -> AccountRow | None:
         now = datetime.now(UTC).isoformat()

karldns/store/base.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+import sqlite3
+from contextlib import AbstractContextManager
+
+
+class StoreConnectionMixin:
+    def connect(self) -> AbstractContextManager[sqlite3.Connection]:
+        raise NotImplementedError

karldns/store/metrics.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 from datetime import UTC, datetime
 
+from .base import StoreConnectionMixin
 
-class MetricsMixin:
+
+class MetricsMixin(StoreConnectionMixin):
     def increment_metric(self, name: str, amount: float = 1.0) -> None:
         now = datetime.now(UTC).isoformat()
         with self.connect() as conn:

karldns/store/rate_limit.py

@@ -1,14 +1,23 @@
 from __future__ import annotations
 
+import threading
 import time
 from datetime import UTC, datetime, timedelta
 
-from .types import CleanupRunRow
+from .base import StoreConnectionMixin
+from .types import CleanupRunRow, cleanup_run_row
 
 RATE_LIMIT_FLUSH_SECONDS = 10.0
 
 
-class RateLimitCleanupMixin:
+class RateLimitCleanupMixin(StoreConnectionMixin):
+    _rate_limit_lock: threading.Lock
+    _rate_limit_prune_lock: threading.Lock
+    _rate_limit_buckets: dict[str, dict[str, float | int]]
+    _rate_limit_dirty_counts: dict[tuple[str, int], int]
+    _last_rate_limit_flush: float
+    _last_rate_limit_prune_window: int
+
     def list_cleanup_runs(self, limit: int = 20, offset: int = 0) -> list[CleanupRunRow]:
         with self.connect() as conn:
             rows = conn.execute(
@@ -20,7 +29,7 @@ class RateLimitCleanupMixin:
                 """,
                 (limit, offset),
             ).fetchall()
-        return [CleanupRunRow(**dict(row)) for row in rows]
+        return [cleanup_run_row(row) for row in rows]
 
     def count_cleanup_runs(self) -> int:
         with self.connect() as conn:

karldns/store/records.py

@@ -4,10 +4,11 @@ import sqlite3
 from datetime import UTC, datetime
 
 from ..models import UpdateResult
-from .types import UpdateEventRow
+from .base import StoreConnectionMixin
+from .types import RecordRow, UpdateEventRow, record_row, update_event_row
 
 
-class RecordsEventsMixin:
+class RecordsEventsMixin(StoreConnectionMixin):
     def upsert(
         self,
         hostname: str,
@@ -84,10 +85,10 @@ class RecordsEventsMixin:
             update_ipv6=update_ipv6,
         )
 
-    def list_records(self) -> list[dict[str, str | None]]:
+    def list_records(self) -> list[RecordRow]:
         return self.list_records_page(limit=10_000, offset=0)
 
-    def list_records_page(self, limit: int = 100, offset: int = 0) -> list[dict[str, str | None]]:
+    def list_records_page(self, limit: int = 100, offset: int = 0) -> list[RecordRow]:
         with self.connect() as conn:
             rows = conn.execute(
                 """
@@ -98,7 +99,7 @@ class RecordsEventsMixin:
                 """,
                 (limit, offset),
             ).fetchall()
-        return [dict(row) for row in rows]
+        return [record_row(row) for row in rows]
 
     def count_records(self) -> int:
         with self.connect() as conn:
@@ -167,7 +168,7 @@ class RecordsEventsMixin:
                 """,
                 (limit, offset),
             ).fetchall()
-        return [dict(row) for row in rows]
+        return [update_event_row(row) for row in rows]
 
     def list_update_events_before(self, limit: int = 100, before: str | None = None) -> list[UpdateEventRow]:
         where_clause = "WHERE created_at < ?" if before else ""
@@ -189,7 +190,7 @@ class RecordsEventsMixin:
                 """,
                 params,
             ).fetchall()
-        return [dict(row) for row in rows]
+        return [update_event_row(row) for row in rows]
 
     def count_update_events(self) -> int:
         with self.connect() as conn:
@@ -228,7 +229,7 @@ def _merge_event_rows(
     limit: int,
     offset: int,
 ) -> list[UpdateEventRow]:
-    rows = [dict(row) for row in event_rows]
-    rows.extend(dict(row) for row in rollup_rows)
+    rows = [update_event_row(row) for row in event_rows]
+    rows.extend(update_event_row(row) for row in rollup_rows)
     rows.sort(key=lambda row: str(row["created_at"]), reverse=True)
     return rows[offset : offset + limit]

karldns/store/scheduler.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 from datetime import UTC, datetime, timedelta
 
-from .types import ScheduledTaskRow
+from .base import StoreConnectionMixin
+from .types import ScheduledTaskRow, scheduled_task_row
 
 
 def _now() -> datetime:
@@ -13,7 +14,7 @@ def _iso(value: datetime) -> str:
     return value.isoformat()
 
 
-class SchedulerMixin:
+class SchedulerMixin(StoreConnectionMixin):
     def ensure_scheduled_task(self, name: str, interval_seconds: int) -> None:
         now = _now()
         with self.connect() as conn:
@@ -86,12 +87,10 @@ class SchedulerMixin:
 
     def list_scheduled_tasks(self) -> list[ScheduledTaskRow]:
         with self.connect() as conn:
-            rows = conn.execute(
-                """
+            rows = conn.execute("""
                 SELECT name, interval_seconds, next_run_at, lease_owner, lease_expires_at,
                        last_started_at, last_finished_at, last_status, last_error, updated_at
                 FROM scheduled_tasks
                 ORDER BY name
-                """
-            ).fetchall()
-        return [ScheduledTaskRow(**dict(row)) for row in rows]
+                """).fetchall()
+        return [scheduled_task_row(row) for row in rows]

karldns/store/schema.py

@@ -35,18 +35,15 @@ def initialize_schema(conn: sqlite3.Connection) -> None:
     configure_connection(conn)
     schema_version = conn.execute("PRAGMA user_version").fetchone()[0]
     conn.execute("PRAGMA journal_mode = WAL")
-    conn.execute(
-        """
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS records (
           hostname TEXT PRIMARY KEY,
           ipv4 TEXT,
           ipv6 TEXT,
           updated_at TEXT NOT NULL
         )
-        """
-    )
-    conn.execute(
-        """
+        """)
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS update_events (
           id INTEGER PRIMARY KEY AUTOINCREMENT,
           hostname TEXT NOT NULL,
@@ -57,12 +54,12 @@ def initialize_schema(conn: sqlite3.Connection) -> None:
           source_ip TEXT NOT NULL,
           created_at TEXT NOT NULL
         )
-        """
-    )
+        """)
     conn.execute("CREATE INDEX IF NOT EXISTS idx_update_events_created_at ON update_events(created_at)")
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_update_events_hostname_created_at ON update_events(hostname, created_at DESC)")
     conn.execute(
-        """
+        "CREATE INDEX IF NOT EXISTS idx_update_events_hostname_created_at ON update_events(hostname, created_at DESC)"
+    )
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS update_event_rollups (
           id INTEGER PRIMARY KEY AUTOINCREMENT,
           hostname TEXT NOT NULL,
@@ -76,28 +73,24 @@ def initialize_schema(conn: sqlite3.Connection) -> None:
           count INTEGER NOT NULL,
           UNIQUE(hostname, ipv4, ipv6, status, detail, source_ip)
         )
-        """
-    )
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_update_event_rollups_last_seen_at ON update_event_rollups(last_seen_at)")
+        """)
     conn.execute(
-        """
+        "CREATE INDEX IF NOT EXISTS idx_update_event_rollups_last_seen_at ON update_event_rollups(last_seen_at)"
+    )
+    conn.execute("""
         CREATE INDEX IF NOT EXISTS idx_update_event_rollups_hostname_last_seen_at
         ON update_event_rollups(hostname, last_seen_at DESC)
-        """
-    )
-    conn.execute(
-        """
+        """)
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS cleanup_runs (
           id INTEGER PRIMARY KEY AUTOINCREMENT,
           created_at TEXT NOT NULL,
           unused_accounts_deleted INTEGER NOT NULL
         )
-        """
-    )
+        """)
     _migrate_cleanup_runs_schema(conn)
     conn.execute("CREATE INDEX IF NOT EXISTS idx_cleanup_runs_created_at ON cleanup_runs(created_at)")
-    conn.execute(
-        """
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS scheduled_tasks (
           name TEXT PRIMARY KEY,
           interval_seconds INTEGER NOT NULL,
@@ -110,37 +103,29 @@ def initialize_schema(conn: sqlite3.Connection) -> None:
           last_error TEXT NOT NULL DEFAULT '',
           updated_at TEXT NOT NULL
         )
-        """
-    )
+        """)
     conn.execute("CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_next_run_at ON scheduled_tasks(next_run_at)")
-    conn.execute(
-        """
+    conn.execute("""
         CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_due
         ON scheduled_tasks(next_run_at, lease_expires_at)
-        """
-    )
-    conn.execute(
-        """
+        """)
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS rate_limits (
           key TEXT NOT NULL,
           window INTEGER NOT NULL,
           count INTEGER NOT NULL,
           PRIMARY KEY(key, window)
         )
-        """
-    )
+        """)
     conn.execute("CREATE INDEX IF NOT EXISTS idx_rate_limits_window ON rate_limits(window)")
-    conn.execute(
-        """
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS service_metrics (
           name TEXT PRIMARY KEY,
           value REAL NOT NULL,
           updated_at TEXT NOT NULL
         )
-        """
-    )
-    conn.execute(
-        """
+        """)
+    conn.execute("""
         CREATE TABLE IF NOT EXISTS accounts (
           hostname TEXT PRIMARY KEY,
           update_slug TEXT NOT NULL UNIQUE,
@@ -150,8 +135,7 @@ def initialize_schema(conn: sqlite3.Connection) -> None:
           accepted_at TEXT,
           disabled INTEGER NOT NULL DEFAULT 0
         )
-        """
-    )
+        """)
     _ensure_column(conn, "accounts", "update_slug", "TEXT")
     _ensure_column(conn, "accounts", "management_hash", "TEXT")
     _ensure_column(conn, "accounts", "dashboard_password_hash", "TEXT")
@@ -190,8 +174,7 @@ def _migrate_accounts_schema(conn: sqlite3.Connection) -> None:
     conn.execute("DROP INDEX IF EXISTS idx_accounts_accepted_at")
     conn.execute("DROP INDEX IF EXISTS idx_accounts_created_at")
     conn.execute("DROP INDEX IF EXISTS idx_accounts_hostname_username")
-    conn.execute(
-        """
+    conn.execute("""
         CREATE TABLE accounts_new (
           hostname TEXT PRIMARY KEY,
           update_slug TEXT NOT NULL UNIQUE,
@@ -201,15 +184,12 @@ def _migrate_accounts_schema(conn: sqlite3.Connection) -> None:
           accepted_at TEXT,
           disabled INTEGER NOT NULL DEFAULT 0
         )
-        """
-    )
-    conn.execute(
-        """
+        """)
+    conn.execute("""
         INSERT INTO accounts_new(hostname, update_slug, management_hash, dashboard_password_hash, created_at, accepted_at, disabled)
         SELECT hostname, update_slug, management_hash, dashboard_password_hash, created_at, accepted_at, disabled
         FROM accounts
-        """
-    )
+        """)
     conn.execute("DROP TABLE accounts")
     conn.execute("ALTER TABLE accounts_new RENAME TO accounts")
 
@@ -219,21 +199,17 @@ def _migrate_cleanup_runs_schema(conn: sqlite3.Connection) -> None:
     if columns <= {"id", "created_at", "unused_accounts_deleted"}:
         return
     conn.execute("DROP INDEX IF EXISTS idx_cleanup_runs_created_at")
-    conn.execute(
-        """
+    conn.execute("""
         CREATE TABLE cleanup_runs_new (
           id INTEGER PRIMARY KEY AUTOINCREMENT,
           created_at TEXT NOT NULL,
           unused_accounts_deleted INTEGER NOT NULL
         )
-        """
-    )
-    conn.execute(
-        """
+        """)
+    conn.execute("""
         INSERT INTO cleanup_runs_new(id, created_at, unused_accounts_deleted)
         SELECT id, created_at, unused_accounts_deleted
         FROM cleanup_runs
-        """
-    )
+        """)
     conn.execute("DROP TABLE cleanup_runs")
     conn.execute("ALTER TABLE cleanup_runs_new RENAME TO cleanup_runs")

karldns/store/types.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TypedDict
+from typing import Any, TypedDict, cast
 
 
 class AccountRow(TypedDict, total=False):
@@ -43,3 +43,30 @@ class ScheduledTaskRow(TypedDict, total=False):
     last_status: str
     last_error: str
     updated_at: str
+
+
+class RecordRow(TypedDict):
+    hostname: str
+    ipv4: str | None
+    ipv6: str | None
+    updated_at: str
+
+
+def account_row(row: Any) -> AccountRow:
+    return cast(AccountRow, dict(row))
+
+
+def cleanup_run_row(row: Any) -> CleanupRunRow:
+    return cast(CleanupRunRow, dict(row))
+
+
+def record_row(row: Any) -> RecordRow:
+    return cast(RecordRow, dict(row))
+
+
+def scheduled_task_row(row: Any) -> ScheduledTaskRow:
+    return cast(ScheduledTaskRow, dict(row))
+
+
+def update_event_row(row: Any) -> UpdateEventRow:
+    return cast(UpdateEventRow, dict(row))