Files
devdisc/db.py
2026-06-08 06:18:44 +00:00

965 lines
25 KiB
Python

from __future__ import annotations
import json
import logging
import queue
import sqlite3
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Any
from config import (
NETWORK_BATCH_SIZE,
NETWORK_FLUSH_INTERVAL,
NETWORK_RETENTION,
)
log = logging.getLogger(__name__)
class Database:
def __init__(
self,
db_path: str | Path,
):
self.db_path = Path(
db_path
)
self.db_path.parent.mkdir(
parents=True,
exist_ok=True,
)
self._lock = threading.RLock()
self._queue: queue.Queue = (
queue.Queue(
maxsize=20000
)
)
self._running = True
self._conn = sqlite3.connect(
self.db_path,
check_same_thread=False,
)
self._conn.row_factory = (
sqlite3.Row
)
self._initialize()
self._writer = threading.Thread(
target=self._writer_loop,
daemon=True,
name="DatabaseWriter",
)
self._writer.start()
@contextmanager
def connection(self):
with self._lock:
yield self._conn
def _initialize(
self,
) -> None:
with self.connection() as conn:
conn.executescript(
"""
PRAGMA journal_mode=WAL;
PRAGMA synchronous=NORMAL;
PRAGMA foreign_keys=ON;
PRAGMA temp_store=MEMORY;
PRAGMA cache_size=-20000;
PRAGMA mmap_size=268435456;
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS profiles (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT UNIQUE NOT NULL,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS themes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
profile_name TEXT NOT NULL,
name TEXT NOT NULL,
path TEXT NOT NULL,
enabled INTEGER NOT NULL DEFAULT 1,
UNIQUE(profile_name,name)
);
CREATE TABLE IF NOT EXISTS extensions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
profile_name TEXT NOT NULL,
name TEXT NOT NULL,
path TEXT NOT NULL,
enabled INTEGER NOT NULL DEFAULT 1,
UNIQUE(profile_name,name)
);
CREATE TABLE IF NOT EXISTS proxies (
profile_name TEXT PRIMARY KEY,
enabled INTEGER NOT NULL DEFAULT 0,
proxy_type TEXT,
host TEXT,
port INTEGER,
username TEXT,
password TEXT
);
CREATE TABLE IF NOT EXISTS blocked_domains (
id INTEGER PRIMARY KEY AUTOINCREMENT,
domain TEXT UNIQUE NOT NULL
);
CREATE TABLE IF NOT EXISTS network_requests (
id INTEGER PRIMARY KEY AUTOINCREMENT,
request_id TEXT,
timestamp TEXT,
method TEXT,
url TEXT,
host TEXT,
status_code INTEGER,
resource_type TEXT,
request_headers TEXT,
response_headers TEXT,
request_body TEXT,
response_body TEXT,
duration_ms REAL
);
CREATE TABLE IF NOT EXISTS plugin_registry (
id INTEGER PRIMARY KEY AUTOINCREMENT,
profile_name TEXT NOT NULL,
plugin_name TEXT NOT NULL,
version TEXT,
enabled INTEGER NOT NULL DEFAULT 1,
installed_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
UNIQUE(profile_name,plugin_name)
);
CREATE TABLE IF NOT EXISTS plugin_settings (
id INTEGER PRIMARY KEY AUTOINCREMENT,
profile_name TEXT NOT NULL,
plugin_name TEXT NOT NULL,
setting_key TEXT NOT NULL,
setting_value TEXT NOT NULL,
UNIQUE(
profile_name,
plugin_name,
setting_key
)
);
CREATE TABLE IF NOT EXISTS plugin_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
profile_name TEXT NOT NULL,
plugin_name TEXT NOT NULL,
level TEXT NOT NULL,
message TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS plugin_crashes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
profile_name TEXT NOT NULL,
plugin_name TEXT NOT NULL,
error TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_req_host
ON network_requests(host);
CREATE INDEX IF NOT EXISTS idx_req_url
ON network_requests(url);
CREATE INDEX IF NOT EXISTS idx_req_status
ON network_requests(status_code);
CREATE INDEX IF NOT EXISTS idx_req_timestamp
ON network_requests(timestamp);
CREATE INDEX IF NOT EXISTS idx_plugin_registry_profile
ON plugin_registry(profile_name);
CREATE INDEX IF NOT EXISTS idx_plugin_settings_lookup
ON plugin_settings(
profile_name,
plugin_name,
setting_key
);
CREATE INDEX IF NOT EXISTS idx_plugin_logs_profile
ON plugin_logs(profile_name);
CREATE INDEX IF NOT EXISTS idx_plugin_crashes_profile
ON plugin_crashes(profile_name);
"""
)
conn.commit()
def shutdown(
self,
) -> None:
self._running = False
if (
self._writer
and self._writer.is_alive()
):
self._writer.join(
timeout=10
)
with self._lock:
try:
self._conn.commit()
self._conn.close()
except Exception:
pass
def _writer_loop(
self,
) -> None:
batch = []
last_flush = (
time.monotonic()
)
while (
self._running
or not self._queue.empty()
):
try:
batch.append(
self._queue.get(
timeout=0.25
)
)
except queue.Empty:
pass
now = time.monotonic()
if (
len(batch)
>= NETWORK_BATCH_SIZE
or (
batch
and (
now
- last_flush
)
>= NETWORK_FLUSH_INTERVAL
)
):
self._flush(batch)
batch.clear()
last_flush = now
if batch:
self._flush(batch)
def _flush(
self,
batch: list,
) -> None:
try:
with self.connection() as conn:
conn.executemany(
"""
INSERT INTO network_requests (
request_id,
timestamp,
method,
url,
host,
status_code,
resource_type,
request_headers,
response_headers,
request_body,
response_body,
duration_ms
)
VALUES (
?,?,?,?,?,?,
?,?,?,?,?,?
)
""",
batch,
)
conn.commit()
self._cleanup()
except Exception:
log.exception(
"db flush failed"
)
def _cleanup(
self,
) -> None:
try:
with self.connection() as conn:
row = conn.execute(
"""
SELECT MAX(id)
FROM network_requests
"""
).fetchone()
max_id = (
row[0]
if row
else None
)
if (
max_id is None
or max_id
<= NETWORK_RETENTION
):
return
conn.execute(
"""
DELETE FROM network_requests
WHERE id < ?
""",
(
max_id
- NETWORK_RETENTION,
),
)
conn.commit()
except Exception:
log.exception(
"cleanup failed"
)
def queue_request(
self,
*,
request_id: str,
timestamp: str,
method: str | None,
url: str | None,
host: str | None,
status_code: int | None,
resource_type: str | None,
request_headers: dict | None,
response_headers: dict | None,
request_body: str | None,
response_body: str | None,
duration_ms: float | None,
) -> None:
try:
self._queue.put_nowait(
(
request_id,
timestamp,
method,
url,
host,
status_code,
resource_type,
json.dumps(
request_headers
or {}
),
json.dumps(
response_headers
or {}
),
request_body,
response_body,
duration_ms,
)
)
except queue.Full:
log.warning(
"network queue full"
)
def get_setting(
self,
key: str,
default: Any = None,
) -> Any:
with self.connection() as conn:
row = conn.execute(
"""
SELECT value
FROM settings
WHERE key=?
""",
(key,),
).fetchone()
if not row:
return default
try:
return json.loads(
row["value"]
)
except Exception:
return default
def set_setting(
self,
key: str,
value: Any,
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT INTO settings(
key,
value
)
VALUES(?,?)
ON CONFLICT(key)
DO UPDATE SET
value=excluded.value
""",
(
key,
json.dumps(
value,
ensure_ascii=False,
),
),
)
conn.commit()
def ensure_profile(
self,
profile_name: str,
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT OR IGNORE
INTO profiles(name)
VALUES(?)
""",
(
profile_name,
),
)
conn.commit()
def list_profiles(
self,
) -> list[str]:
with self.connection() as conn:
rows = conn.execute(
"""
SELECT name
FROM profiles
ORDER BY name
"""
).fetchall()
return [
row["name"]
for row in rows
]
def add_blocked_domain(
self,
domain: str,
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT OR IGNORE
INTO blocked_domains(domain)
VALUES(?)
""",
(
domain.lower().strip(),
),
)
conn.commit()
def blocked_domains(
self,
) -> list[str]:
with self.connection() as conn:
rows = conn.execute(
"""
SELECT domain
FROM blocked_domains
ORDER BY domain
"""
).fetchall()
return [
row["domain"]
for row in rows
]
def recent_requests(
self,
limit: int = 1000,
):
with self.connection() as conn:
return conn.execute(
"""
SELECT *
FROM network_requests
ORDER BY id DESC
LIMIT ?
""",
(
limit,
),
).fetchall()
def search_requests(
self,
query: str,
limit: int = 1000,
):
like = f"%{query}%"
with self.connection() as conn:
return conn.execute(
"""
SELECT *
FROM network_requests
WHERE
url LIKE ?
OR host LIKE ?
OR method LIKE ?
ORDER BY id DESC
LIMIT ?
""",
(
like,
like,
like,
limit,
),
).fetchall()
def plugin_enabled(
self,
profile_name: str,
plugin_name: str,
default: bool = True,
) -> bool:
with self.connection() as conn:
row = conn.execute(
"""
SELECT enabled
FROM plugin_registry
WHERE
profile_name=?
AND plugin_name=?
""",
(
profile_name,
plugin_name,
),
).fetchone()
return (
default
if not row
else bool(
row["enabled"]
)
)
def set_plugin_enabled(
self,
profile_name: str,
plugin_name: str,
enabled: bool,
version: str = "",
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT INTO plugin_registry(
profile_name,
plugin_name,
version,
enabled
)
VALUES(?,?,?,?)
ON CONFLICT(
profile_name,
plugin_name
)
DO UPDATE SET
enabled=excluded.enabled,
version=excluded.version
""",
(
profile_name,
plugin_name,
version,
int(enabled),
),
)
conn.commit()
def plugin_setting_get(
self,
profile_name: str,
plugin_name: str,
key: str,
default=None,
):
with self.connection() as conn:
row = conn.execute(
"""
SELECT setting_value
FROM plugin_settings
WHERE
profile_name=?
AND plugin_name=?
AND setting_key=?
""",
(
profile_name,
plugin_name,
key,
),
).fetchone()
if not row:
return default
try:
return json.loads(
row["setting_value"]
)
except Exception:
return default
def plugin_setting_set(
self,
profile_name: str,
plugin_name: str,
key: str,
value,
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT INTO plugin_settings(
profile_name,
plugin_name,
setting_key,
setting_value
)
VALUES(?,?,?,?)
ON CONFLICT(
profile_name,
plugin_name,
setting_key
)
DO UPDATE SET
setting_value=excluded.setting_value
""",
(
profile_name,
plugin_name,
key,
json.dumps(
value,
ensure_ascii=False,
),
),
)
conn.commit()
def plugin_log(
self,
profile_name: str,
plugin_name: str,
level: str,
message: str,
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT INTO plugin_logs(
profile_name,
plugin_name,
level,
message
)
VALUES(?,?,?,?)
""",
(
profile_name,
plugin_name,
level,
message,
),
)
conn.commit()
def plugin_crash(
self,
profile_name: str,
plugin_name: str,
error: str,
) -> None:
with self.connection() as conn:
conn.execute(
"""
INSERT INTO plugin_crashes(
profile_name,
plugin_name,
error
)
VALUES(?,?,?)
""",
(
profile_name,
plugin_name,
error,
),
)
conn.commit()
def plugin_logs(
self,
profile_name: str,
plugin_name: str | None = None,
limit: int = 1000,
):
with self.connection() as conn:
if plugin_name:
return conn.execute(
"""
SELECT *
FROM plugin_logs
WHERE
profile_name=?
AND plugin_name=?
ORDER BY id DESC
LIMIT ?
""",
(
profile_name,
plugin_name,
limit,
),
).fetchall()
return conn.execute(
"""
SELECT *
FROM plugin_logs
WHERE profile_name=?
ORDER BY id DESC
LIMIT ?
""",
(
profile_name,
limit,
),
).fetchall()
def plugin_crashes(
self,
profile_name: str,
plugin_name: str | None = None,
limit: int = 1000,
):
with self.connection() as conn:
if plugin_name:
return conn.execute(
"""
SELECT *
FROM plugin_crashes
WHERE
profile_name=?
AND plugin_name=?
ORDER BY id DESC
LIMIT ?
""",
(
profile_name,
plugin_name,
limit,
),
).fetchall()
return conn.execute(
"""
SELECT *
FROM plugin_crashes
WHERE profile_name=?
ORDER BY id DESC
LIMIT ?
""",
(
profile_name,
limit,
),
).fetchall()
def plugins(
self,
profile_name: str,
):
with self.connection() as conn:
return conn.execute(
"""
SELECT *
FROM plugin_registry
WHERE profile_name=?
ORDER BY plugin_name
""",
(
profile_name,
),
).fetchall()
# db.py additions
def delete_plugin(
self,
profile_name: str,
plugin_name: str,
) -> None:
with self.connection() as conn:
conn.execute(
"""
DELETE FROM plugin_registry
WHERE
profile_name=?
AND plugin_name=?
""",
(
profile_name,
plugin_name,
),
)
conn.execute(
"""
DELETE FROM plugin_settings
WHERE
profile_name=?
AND plugin_name=?
""",
(
profile_name,
plugin_name,
),
)
conn.execute(
"""
DELETE FROM plugin_logs
WHERE
profile_name=?
AND plugin_name=?
""",
(
profile_name,
plugin_name,
),
)
conn.execute(
"""
DELETE FROM plugin_crashes
WHERE
profile_name=?
AND plugin_name=?
""",
(
profile_name,
plugin_name,
),
)
conn.commit()
# v1
def plugin_crash_count(
self,
profile_name: str,
plugin_name: str,
) -> int:
with self.connection() as conn:
row = conn.execute(
"""
SELECT COUNT(*)
FROM plugin_crashes
WHERE
profile_name=?
AND plugin_name=?
""",
(
profile_name,
plugin_name,
),
).fetchone()
return int(row[0] or 0)
# v1
def plugin_quarantined(
self,
profile_name: str,
plugin_name: str,
) -> bool:
return bool(
self.get_setting(
f"plugin_quarantine:{profile_name}:{plugin_name}",
False,
)
)
# v1
def set_plugin_quarantined(
self,
profile_name: str,
plugin_name: str,
value: bool,
) -> None:
self.set_setting(
f"plugin_quarantine:{profile_name}:{plugin_name}",
bool(value),
)