Initial commit

solar_relay/storage.py

@@ -0,0 +1,263 @@
+"""Durable, date-partitioned Parquet archive and acknowledged upload queue.
+
+Each mutation replaces a daily partition atomically. A process/thread lock serializes
+readers and writers; it is never held during an HTTP request.
+"""
+
+from contextlib import contextmanager
+from datetime import datetime, timedelta, timezone
+import fcntl
+import json
+import os
+from pathlib import Path
+import tempfile
+import threading
+import time
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from .config import utc_text
+from .coverage import daily_coverage, milliseconds, sample_coverage
+
+NUMBERS = (
+    "production_w", "production_ac_w", "consumption_w", "grid_w",
+    "grid_import_w", "grid_export_w", "production_ac_total_wh",
+    "grid_import_total_wh", "grid_export_total_wh",
+)
+SCHEMA = pa.schema([
+    ("id", pa.string()), ("source_id", pa.string()), ("observed_at", pa.string()),
+    ("observed_ms", pa.int64()), ("recorded_ns", pa.int64()),
+    ("uploaded_at", pa.string()), ("payload", pa.string()),
+    *((field, pa.float64()) for field in NUMBERS),
+    ("boiler_enabled", pa.bool_()), ("boiler_reason", pa.string()),
+    ("record_kind", pa.string()),
+    ("interval_energy_wh", pa.struct([(field, pa.float64()) for field in ("production", "consumption", "self_consumption", "bought", "sold")])),
+    ("power_ms", pa.int64()), ("power_ok", pa.bool_()),
+    ("meter_ms", pa.int64()), ("meter_ok", pa.bool_()),
+    ("boiler_ms", pa.int64()), ("boiler_ok", pa.bool_()),
+], metadata={b"solar_relay_schema": b"1"})
+COVERAGE_COLUMNS = ["observed_ms", "power_ms", "power_ok", "meter_ms", "meter_ok", "boiler_ms", "boiler_ok"]
+
+
+class StoreError(OSError):
+    """An archive could not be read or updated safely."""
+
+
+class Store:
+    def __init__(self, path: Path, retention_days=30):
+        self.path = Path(path)
+        if self.path.is_file():
+            raise StoreError("collector.archive must be a directory")
+        self.path.mkdir(parents=True, exist_ok=True, mode=0o700)
+        self.retention_days = retention_days
+        self._thread_lock = threading.RLock()
+        descriptor = os.open(self.path / ".lock", os.O_CREAT | os.O_WRONLY, 0o600)
+        os.close(descriptor)
+
+    @contextmanager
+    def _locked(self, *, write=False):
+        with self._thread_lock, (self.path / ".lock").open("rb") as handle:
+            fcntl.flock(handle, fcntl.LOCK_EX if write else fcntl.LOCK_SH)
+            try:
+                yield
+            except pa.ArrowException as error:
+                raise StoreError(f"Parquet archive error ({type(error).__name__})") from error
+            finally:
+                fcntl.flock(handle, fcntl.LOCK_UN)
+
+    def _files(self):
+        return sorted(self.path.glob("date=????-??-??/samples.parquet"))
+
+    def _partition(self, stamp):
+        date = datetime.fromtimestamp(milliseconds(stamp) / 1000, timezone.utc).date()
+        return self.path / f"date={date.isoformat()}" / "samples.parquet"
+
+    @staticmethod
+    def _read(path, columns=None):
+        # Avoid implicit Hive partition/schema inference from the directory name.
+        if not path.exists():
+            return []
+        parquet = pq.ParquetFile(path, page_checksum_verification=True)
+        if not parquet.schema_arrow.equals(SCHEMA, check_metadata=True):
+            raise StoreError("Unsupported Parquet archive schema")
+        return parquet.read(columns=columns).to_pylist()
+
+    @staticmethod
+    def _sync_directory(path):
+        descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
+        try:
+            os.fsync(descriptor)
+        finally:
+            os.close(descriptor)
+
+    def _write(self, path, rows):
+        new_directory = not path.parent.exists()
+        path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+        if new_directory:
+            self._sync_directory(self.path)
+        if not rows:
+            path.unlink(missing_ok=True)
+            self._sync_directory(path.parent)
+            return
+        descriptor, temporary = tempfile.mkstemp(prefix=".samples-", suffix=".tmp", dir=path.parent)
+        try:
+            with os.fdopen(descriptor, "wb") as handle:
+                table = pa.Table.from_pylist(sorted(rows, key=lambda row: (row["observed_ms"], row["recorded_ns"], row["id"])), schema=SCHEMA)
+                pq.write_table(table, handle, compression="zstd", write_page_checksum=True)
+                handle.flush()
+                os.fsync(handle.fileno())
+            os.replace(temporary, path)
+            self._sync_directory(path.parent)
+        finally:
+            Path(temporary).unlink(missing_ok=True)
+
+    @staticmethod
+    def _row(sample, uploaded_at=None):
+        coverage = sample_coverage(sample)
+        readings = sample.get("readings", {})
+        return {
+            "id": sample["id"], "source_id": sample["source_id"], "observed_at": sample["observed_at"],
+            "observed_ms": coverage[2], "recorded_ns": time.time_ns(), "uploaded_at": uploaded_at,
+            "payload": json.dumps(sample, allow_nan=False, separators=(",", ":")),
+            **{name: readings.get(name) for name in NUMBERS},
+            "boiler_enabled": readings.get("boiler_enabled"), "boiler_reason": readings.get("boiler_reason"),
+            "record_kind": "history" if "history_interval" in sample else "live",
+            "interval_energy_wh": sample.get("history_interval", {}).get("energy_wh"),
+            **{name: bool(value) if name.endswith("_ok") else value
+               for name, value in zip(COVERAGE_COLUMNS[1:], coverage[3:])},
+        }
+
+    def add(self, sample):
+        path, row = self._partition(sample["observed_at"]), self._row(sample)
+        with self._locked(write=True):
+            rows = self._read(path)
+            for existing in rows:
+                if existing["id"] == row["id"]:
+                    if existing["payload"] != row["payload"]:
+                        raise StoreError("Sample ID already contains different data")
+                    return
+            self._write(path, [*rows, row])
+
+    def pending(self, limit):
+        with self._locked():
+            rows = []
+            for path in self._files():
+                index = self._read(path, ["uploaded_at"])
+                if any(row["uploaded_at"] is None for row in index):
+                    rows.extend(row for row in self._read(path, ["recorded_ns", "id", "payload", "uploaded_at"]) if row["uploaded_at"] is None)
+            rows.sort(key=lambda row: (row["recorded_ns"], row["id"]))
+            return [json.loads(row["payload"]) for row in rows[:limit]]
+
+    def acknowledge(self, ids, *, now=None):
+        wanted, stamp = set(ids), utc_text(now or datetime.now(timezone.utc))
+        if not wanted:
+            return
+        with self._locked(write=True):
+            for path in self._files():
+                index = self._read(path, ["id", "uploaded_at"])
+                if not any(row["id"] in wanted and row["uploaded_at"] is None for row in index):
+                    continue
+                rows = self._read(path)
+                for row in rows:
+                    if row["id"] in wanted and row["uploaded_at"] is None:
+                        row["uploaded_at"] = stamp
+                self._write(path, rows)
+
+    def records(self, start=None, until=None):
+        lower, upper = milliseconds(start) if start else None, milliseconds(until) if until else None
+        with self._locked():
+            rows = []
+            for path in self._files():
+                rows.extend(row for row in self._read(path, ["observed_ms", "recorded_ns", "id", "payload"])
+                            if (lower is None or row["observed_ms"] >= lower) and (upper is None or row["observed_ms"] < upper))
+        rows.sort(key=lambda row: (row["observed_ms"], row["recorded_ns"], row["id"]))
+        return (json.loads(row["payload"]) for row in rows)
+
+    def replay(self, start=None, until=None):
+        lower, upper = milliseconds(start) if start else None, milliseconds(until) if until else None
+        count = 0
+        with self._locked(write=True):
+            for path in self._files():
+                rows, changed = self._read(path), False
+                for row in rows:
+                    if (lower is None or row["observed_ms"] >= lower) and (upper is None or row["observed_ms"] < upper):
+                        row["uploaded_at"] = None
+                        count += 1
+                        changed = True
+                if changed:
+                    self._write(path, rows)
+        return count
+
+    def prune(self, *, now=None):
+        """Retain pending records; age acknowledged records from their receipt time."""
+        cutoff = utc_text((now or datetime.now(timezone.utc)) - timedelta(days=self.retention_days))
+        deleted = 0
+        with self._locked(write=True):
+            for path in self._files():
+                index = self._read(path, ["uploaded_at"])
+                if not any(row["uploaded_at"] is not None and row["uploaded_at"] < cutoff for row in index):
+                    continue
+                rows = self._read(path)
+                retained = [row for row in rows if row["uploaded_at"] is None or row["uploaded_at"] >= cutoff]
+                self._write(path, retained)
+                deleted += len(rows) - len(retained)
+        return deleted
+
+    def status(self):
+        with self._locked():
+            rows = []
+            latest = None
+            for path in self._files():
+                index = self._read(path, ["observed_at", "observed_ms", "recorded_ns", "uploaded_at"])
+                rows.extend(index)
+                if index:
+                    candidate = max(index, key=lambda row: (row["observed_ms"], row["recorded_ns"]))
+                    if latest is None or (candidate["observed_ms"], candidate["recorded_ns"]) > latest[0]:
+                        latest = ((candidate["observed_ms"], candidate["recorded_ns"]), path)
+            errors = {}
+            if latest:
+                newest = max(self._read(latest[1], ["observed_ms", "recorded_ns", "payload"]), key=lambda row: (row["observed_ms"], row["recorded_ns"]))
+                errors = json.loads(newest["payload"]).get("errors", {})
+        return {
+            "archive": str(self.path), "format": "parquet", "retention_days": self.retention_days,
+            "samples": len(rows), "pending": sum(row["uploaded_at"] is None for row in rows),
+            "first_observation": min((row["observed_at"] for row in rows), default=None),
+            "last_observation": max((row["observed_at"] for row in rows), default=None),
+            "last_upload": max((row["uploaded_at"] for row in rows if row["uploaded_at"]), default=None),
+            "last_errors": errors,
+        }
+
+    def coverage(self, source_id, timezone_name, days, now=None):
+        from zoneinfo import ZoneInfo
+        now = now or datetime.now(timezone.utc)
+        first = now.astimezone(ZoneInfo(timezone_name)).date() - timedelta(days=days)
+        lower = round(datetime.combine(first, datetime.min.time(), ZoneInfo(timezone_name)).timestamp() * 1000)
+        upper = round(now.timestamp() * 1000)
+        with self._locked():
+            rows = []
+            for path in self._files():
+                if path.parent.name[5:] < (now - timedelta(days=days + 2)).date().isoformat():
+                    continue
+                rows.extend(tuple(row[name] for name in COVERAGE_COLUMNS)
+                            for row in self._read(path, ["source_id", *COVERAGE_COLUMNS])
+                            if row["source_id"] == source_id and lower <= row["observed_ms"] <= upper)
+        return daily_coverage(sorted(rows), timezone_name, days, now)
+
+    def last_success(self, source_id):
+        with self._locked():
+            for path in reversed(self._files()):
+                rows = self._read(path, ["source_id", "observed_ms", "power_ok", "meter_ok", "boiler_ok"])
+                stamps = [row["observed_ms"] for row in rows if row["source_id"] == source_id and row["power_ok"] and row["meter_ok"] and row["boiler_ok"]]
+                if stamps:
+                    return utc_text(datetime.fromtimestamp(max(stamps) / 1000, timezone.utc))
+        return None
+
+    def energy_slots(self, source_id, start, until):
+        """Return UTC bucket starts already having valid local or recovered energy."""
+        lower, upper = milliseconds(start), milliseconds(until)
+        with self._locked():
+            return {row["observed_ms"] // 300_000 * 300
+                    for path in self._files()
+                    for row in self._read(path, ["source_id", "observed_ms", "power_ok"])
+                    if row["source_id"] == source_id and row["power_ok"] and lower <= row["observed_ms"] < upper}