99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
 
|
"""The pasture's containers from the workstation: commands, Postgres rows as JSON, Mongo documents as JSON, a copy of
|
||
|
|
GoToSocial's sqlite. Only reads go through here; nothing a peer does is faked in its database."""
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import sqlite3
|
||
|
|
import subprocess
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
|
||
|
|
def run(*args, input=None, check=True, timeout=300):
|
||
|
|
proc = subprocess.run(["podman", *args], input=input, capture_output=True, timeout=timeout,
|
||
|
|
text=isinstance(input, str) or input is None)
|
||
|
|
if check and proc.returncode != 0:
|
||
|
|
raise RuntimeError(f"podman {' '.join(args[:3])}… failed: {proc.stderr.strip()[:500]}")
|
||
|
|
return proc.stdout
|
||
|
|
|
||
|
|
|
||
|
|
def exec_(container, *cmd, input=None, check=True, timeout=300, workdir=None):
|
||
|
|
args = ["exec"]
|
||
|
|
if input is not None:
|
||
|
|
args.append("-i")
|
||
|
|
if workdir:
|
||
|
|
args += ["-w", workdir]
|
||
|
|
return run(*args, container, *cmd, input=input, check=check, timeout=timeout)
|
||
|
|
|
||
|
|
|
||
|
|
def exists(container):
|
||
|
|
return subprocess.run(["podman", "container", "exists", container]).returncode == 0
|
||
|
|
|
||
|
|
|
||
|
|
def running():
|
||
|
|
out = run("ps", "--format", "{{.Names}}")
|
||
|
|
return [n for n in out.split() if n.startswith("pasture-")]
|
||
|
|
|
||
|
|
|
||
|
|
def psql(db, sql, *params):
|
||
|
|
"""Rows of `sql` as a list of dicts. Parameters are bound by psql (:'p1', :'p2'...), never pasted into the SQL."""
|
||
|
|
args = ["exec", "-i", "pasture-postgres", "psql", "-U", "pasture", "-d", db, "-tAX", "-v", "ON_ERROR_STOP=1"]
|
||
|
|
for i, p in enumerate(params, 1):
|
||
|
|
args += ["-v", f"p{i}={p}"]
|
||
|
|
wrapped = f"select coalesce(json_agg(t), '[]'::json) from ({sql.rstrip().rstrip(';')}) t;\n"
|
||
|
|
out = run(*args, input=wrapped)
|
||
|
|
return json.loads(out.strip() or "[]")
|
||
|
|
|
||
|
|
|
||
|
|
def psql_value(db, sql, *params):
|
||
|
|
rows = psql(db, sql, *params)
|
||
|
|
if not rows:
|
||
|
|
return None
|
||
|
|
return next(iter(rows[0].values()))
|
||
|
|
|
||
|
|
|
||
|
|
def mongo(js, db="PrivaPub"):
|
||
|
|
"""The JSON value of a mongosh expression, e.g. `db.Post.find({...}).toArray()`."""
|
||
|
|
script = f"print(EJSON.stringify(({js}), {{relaxed: true}}))"
|
||
|
|
out = run("exec", "pasture-mongo", "mongosh", "--quiet", db, "--eval", script, timeout=120)
|
||
|
|
return json.loads(out.strip().splitlines()[-1]) if out.strip() else None
|
||
|
|
|
||
|
|
|
||
|
|
class SqliteCopy:
|
||
|
|
"""A consistent copy of a container's sqlite database (with its WAL), opened read-only on the workstation."""
|
||
|
|
|
||
|
|
def __init__(self, container, path):
|
||
|
|
self.dir = tempfile.mkdtemp(prefix="town-sqlite-")
|
||
|
|
local = os.path.join(self.dir, os.path.basename(path))
|
||
|
|
run("cp", f"{container}:{path}", local)
|
||
|
|
for suffix in ("-wal", "-shm"):
|
||
|
|
subprocess.run(["podman", "cp", f"{container}:{path}{suffix}", local + suffix], capture_output=True)
|
||
|
|
self.db = sqlite3.connect(local)
|
||
|
|
self.db.row_factory = sqlite3.Row
|
||
|
|
|
||
|
|
def rows(self, sql, *params):
|
||
|
|
return [dict(r) for r in self.db.execute(sql, params)]
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
self.db.close()
|
||
|
|
shutil.rmtree(self.dir, ignore_errors=True)
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *exc):
|
||
|
|
self.close()
|
||
|
|
|
||
|
|
|
||
|
|
def logs(container, since=None):
|
||
|
|
args = ["logs"]
|
||
|
|
if since:
|
||
|
|
args += ["--since", since]
|
||
|
|
return run(*args, container, check=False)
|
||
|
|
|
||
|
|
|
||
|
|
def stats():
|
||
|
|
out = run("stats", "--no-stream", "--format", "json", *running(), check=False)
|
||
|
|
try:
|
||
|
|
return json.loads(out)
|
||
|
|
except ValueError:
|
||
|
|
return []
|