SocialPub/tools/pasture/town/core/podman.py

"""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 mysql_json(db, sql):
	"""The JSON value one MySQL query returns (built with json_object/json_arrayagg, so text keeps its newlines), from the
	shared MySQL. The mysql client binds nothing: values go in through `mysql_quote`."""
	out = run("exec", "pasture-mysql", "mysql", "-upasture", "-ppasture", "--default-character-set=utf8mb4", "-N", "-B", "--raw", db,
			  "-e", sql)
	text = "\n".join(line for line in out.splitlines() if not line.startswith("mysql: [Warning]")).strip()
	return json.loads(text) if text and text != "NULL" else None


def mysql_quote(value):
	return "'" + str(value).replace("\\", "\\\\").replace("'", "''") + "'"


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 []