264 lines
12 KiB
Python
264 lines
12 KiB
Python
"""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}
|