246 lines
9.7 KiB
Python
246 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
import sqlite3
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def utc_now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
def to_db_time(value: datetime) -> str:
|
|
return value.astimezone(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z")
|
|
|
|
|
|
def from_db_time(value: str) -> datetime:
|
|
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
|
|
|
|
|
class Storage:
|
|
def __init__(self, path: str):
|
|
self.path = path
|
|
self.db: sqlite3.Connection | None = None
|
|
|
|
def open(self) -> None:
|
|
Path(self.path).parent.mkdir(parents=True, exist_ok=True)
|
|
self.db = sqlite3.connect(self.path, check_same_thread=False)
|
|
self.db.row_factory = sqlite3.Row
|
|
self.db.execute("PRAGMA journal_mode=WAL")
|
|
self.db.execute("PRAGMA synchronous=NORMAL")
|
|
self.db.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS samples (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
sampled_at TEXT NOT NULL,
|
|
watts REAL NOT NULL,
|
|
voltage REAL,
|
|
current_amps REAL,
|
|
socket_on INTEGER NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_samples_at ON samples(sampled_at);
|
|
|
|
CREATE TABLE IF NOT EXISTS cycles (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
started_at TEXT NOT NULL,
|
|
ended_at TEXT,
|
|
duration_seconds REAL,
|
|
average_watts REAL,
|
|
maximum_watts REAL,
|
|
energy_wh REAL,
|
|
sample_count INTEGER,
|
|
incomplete_start INTEGER NOT NULL DEFAULT 0,
|
|
incomplete_end INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_cycles_started ON cycles(started_at);
|
|
"""
|
|
)
|
|
self.db.commit()
|
|
|
|
def close(self) -> None:
|
|
if self.db is not None:
|
|
self.db.close()
|
|
self.db = None
|
|
|
|
def _conn(self) -> sqlite3.Connection:
|
|
if self.db is None:
|
|
raise RuntimeError("Database is not open")
|
|
return self.db
|
|
|
|
def record_sample(
|
|
self,
|
|
sampled_at: datetime,
|
|
watts: float,
|
|
voltage: float | None,
|
|
current_amps: float | None,
|
|
socket_on: bool,
|
|
) -> None:
|
|
db = self._conn()
|
|
db.execute(
|
|
"INSERT INTO samples(sampled_at, watts, voltage, current_amps, socket_on) VALUES (?, ?, ?, ?, ?)",
|
|
(to_db_time(sampled_at), watts, voltage, current_amps, int(socket_on)),
|
|
)
|
|
db.commit()
|
|
|
|
def start_cycle(self, started_at: datetime, incomplete_start: bool = False) -> int:
|
|
db = self._conn()
|
|
cursor = db.execute(
|
|
"INSERT INTO cycles(started_at, incomplete_start) VALUES (?, ?)",
|
|
(to_db_time(started_at), int(incomplete_start)),
|
|
)
|
|
db.commit()
|
|
return int(cursor.lastrowid)
|
|
|
|
def finish_cycle(self, cycle_id: int, ended_at: datetime, incomplete_end: bool = False) -> dict[str, Any]:
|
|
db = self._conn()
|
|
cycle = db.execute("SELECT * FROM cycles WHERE id = ?", (cycle_id,)).fetchone()
|
|
if cycle is None:
|
|
raise RuntimeError(f"Cycle {cycle_id} does not exist")
|
|
|
|
started_at = from_db_time(cycle["started_at"])
|
|
if ended_at < started_at:
|
|
ended_at = started_at
|
|
start_text = to_db_time(started_at)
|
|
end_text = to_db_time(ended_at)
|
|
rows = db.execute(
|
|
"SELECT sampled_at, watts FROM samples WHERE sampled_at >= ? AND sampled_at < ? ORDER BY sampled_at",
|
|
(start_text, end_text),
|
|
).fetchall()
|
|
|
|
values = [(from_db_time(row["sampled_at"]), float(row["watts"])) for row in rows]
|
|
duration = max(0.0, (ended_at - started_at).total_seconds())
|
|
average = sum(watts for _, watts in values) / len(values) if values else 0.0
|
|
maximum = max((watts for _, watts in values), default=0.0)
|
|
energy_wh = self._integrate_energy(values, ended_at)
|
|
|
|
db.execute(
|
|
"""
|
|
UPDATE cycles
|
|
SET ended_at = ?, duration_seconds = ?, average_watts = ?, maximum_watts = ?,
|
|
energy_wh = ?, sample_count = ?, incomplete_end = ?
|
|
WHERE id = ?
|
|
""",
|
|
(end_text, duration, average, maximum, energy_wh, len(values), int(incomplete_end), cycle_id),
|
|
)
|
|
db.commit()
|
|
return self.get_cycle(cycle_id)
|
|
|
|
@staticmethod
|
|
def _integrate_energy(values: list[tuple[datetime, float]], ended_at: datetime) -> float:
|
|
if not values:
|
|
return 0.0
|
|
watt_seconds = 0.0
|
|
for index, (sampled_at, watts) in enumerate(values):
|
|
segment_end = values[index + 1][0] if index + 1 < len(values) else ended_at
|
|
seconds = max(0.0, (segment_end - sampled_at).total_seconds())
|
|
watt_seconds += watts * seconds
|
|
return watt_seconds / 3600.0
|
|
|
|
def get_cycle(self, cycle_id: int) -> dict[str, Any]:
|
|
row = self._conn().execute("SELECT * FROM cycles WHERE id = ?", (cycle_id,)).fetchone()
|
|
if row is None:
|
|
raise RuntimeError(f"Cycle {cycle_id} does not exist")
|
|
return dict(row)
|
|
|
|
def close_open_cycles(self) -> int:
|
|
db = self._conn()
|
|
open_rows = db.execute("SELECT id, started_at FROM cycles WHERE ended_at IS NULL").fetchall()
|
|
closed = 0
|
|
for row in open_rows:
|
|
last_sample = db.execute(
|
|
"SELECT sampled_at FROM samples WHERE sampled_at >= ? ORDER BY sampled_at DESC LIMIT 1",
|
|
(row["started_at"],),
|
|
).fetchone()
|
|
ended_at = from_db_time(last_sample["sampled_at"]) if last_sample else from_db_time(row["started_at"])
|
|
self.finish_cycle(int(row["id"]), ended_at, incomplete_end=True)
|
|
closed += 1
|
|
return closed
|
|
|
|
def recent_cycles(self, limit: int = 100) -> list[dict[str, Any]]:
|
|
rows = self._conn().execute(
|
|
"SELECT * FROM cycles WHERE ended_at IS NOT NULL ORDER BY started_at DESC LIMIT ?",
|
|
(max(1, min(limit, 1000)),),
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def cycles_between(self, start: datetime, end: datetime) -> list[dict[str, Any]]:
|
|
rows = self._conn().execute(
|
|
"""
|
|
SELECT * FROM cycles
|
|
WHERE started_at < ? AND COALESCE(ended_at, ?) >= ?
|
|
ORDER BY started_at DESC
|
|
""",
|
|
(to_db_time(end), to_db_time(end), to_db_time(start)),
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def summary_between(self, start: datetime, end: datetime) -> dict[str, Any]:
|
|
cycles = self.cycles_between(start, end)
|
|
durations: list[float] = []
|
|
energy_wh = 0.0
|
|
for cycle in cycles:
|
|
cycle_start = max(start, from_db_time(cycle["started_at"]))
|
|
cycle_end_raw = from_db_time(cycle["ended_at"]) if cycle["ended_at"] else min(end, utc_now())
|
|
cycle_end = min(end, cycle_end_raw)
|
|
durations.append(max(0.0, (cycle_end - cycle_start).total_seconds()))
|
|
energy_wh += float(cycle["energy_wh"] or 0.0)
|
|
return {
|
|
"cycle_count": len(cycles),
|
|
"runtime_seconds": sum(durations),
|
|
"average_duration_seconds": sum(durations) / len(durations) if durations else 0.0,
|
|
"longest_duration_seconds": max(durations, default=0.0),
|
|
"energy_wh": energy_wh,
|
|
}
|
|
|
|
def sample_history(
|
|
self,
|
|
since: datetime,
|
|
until: datetime | None = None,
|
|
max_points: int = 800,
|
|
preserve_ranges: list[tuple[datetime, datetime]] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
until = until or utc_now()
|
|
rows = self._conn().execute(
|
|
"SELECT sampled_at, watts FROM samples WHERE sampled_at >= ? AND sampled_at <= ? ORDER BY sampled_at",
|
|
(to_db_time(since), to_db_time(until)),
|
|
).fetchall()
|
|
|
|
if len(rows) <= max_points:
|
|
selected = list(rows)
|
|
else:
|
|
# First/minimum/maximum/last preserves short power spikes that a
|
|
# simple stride can skip entirely. All readings inside known pump
|
|
# cycles are also retained for accurate hover details.
|
|
bucket_count = max(1, max_points // 4)
|
|
bucket_size = max(1, math.ceil(len(rows) / bucket_count))
|
|
selected_by_time: dict[str, sqlite3.Row] = {}
|
|
for offset in range(0, len(rows), bucket_size):
|
|
bucket = rows[offset : offset + bucket_size]
|
|
for row in (
|
|
bucket[0],
|
|
min(bucket, key=lambda item: float(item["watts"])),
|
|
max(bucket, key=lambda item: float(item["watts"])),
|
|
bucket[-1],
|
|
):
|
|
selected_by_time[row["sampled_at"]] = row
|
|
|
|
ranges = preserve_ranges or []
|
|
if ranges:
|
|
normalized = [(start.astimezone(UTC), end.astimezone(UTC)) for start, end in ranges]
|
|
for row in rows:
|
|
sampled_at = from_db_time(row["sampled_at"])
|
|
if any(start <= sampled_at <= end for start, end in normalized):
|
|
selected_by_time[row["sampled_at"]] = row
|
|
|
|
selected = sorted(selected_by_time.values(), key=lambda row: row["sampled_at"])
|
|
return [{"sampled_at": row["sampled_at"], "watts": row["watts"]} for row in selected]
|
|
|
|
def cleanup_samples(self, retention_days: int) -> int:
|
|
cutoff = to_db_time(utc_now() - timedelta(days=retention_days))
|
|
db = self._conn()
|
|
cursor = db.execute("DELETE FROM samples WHERE sampled_at < ?", (cutoff,))
|
|
db.commit()
|
|
return int(cursor.rowcount)
|