diff --git a/puppet/zulip/files/postgresql/env-wal-g b/puppet/zulip/files/postgresql/env-wal-g index a397e8d637..e68c7604c3 100755 --- a/puppet/zulip/files/postgresql/env-wal-g +++ b/puppet/zulip/files/postgresql/env-wal-g @@ -10,7 +10,7 @@ export WALG_DELTA_MAX_STEPS s3_backups_bucket=$(crudini --get "$ZULIP_SECRETS_CONF" secrets s3_backups_bucket 2>/dev/null) -if [ "$s3_backups_bucket" != "" ]; then +if [ "$s3_backups_bucket" != "" ] || [ -n "$WALG_S3_PREFIX" ]; then AWS_REGION=$(crudini --get "$ZULIP_SECRETS_CONF" secrets s3_region 2>/dev/null) if [ "$AWS_REGION" = "" ]; then # Fall back to the current region, if possible @@ -24,7 +24,8 @@ if [ "$s3_backups_bucket" != "" ]; then export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY=$(crudini --get "$ZULIP_SECRETS_CONF" secrets s3_backups_secret_key 2>/dev/null) export AWS_SECRET_ACCESS_KEY - export WALG_S3_PREFIX="s3://$s3_backups_bucket" + : "${WALG_S3_PREFIX:=s3://$s3_backups_bucket}" + export WALG_S3_PREFIX if storage_class=$(crudini --get /etc/zulip/zulip.conf postgresql backups_storage_class 2>&1); then export WALG_S3_STORAGE_CLASS="$storage_class" diff --git a/puppet/zulip/files/postgresql/wal-g-exporter b/puppet/zulip/files/postgresql/wal-g-exporter index edb2790717..b5fb1bef52 100755 --- a/puppet/zulip/files/postgresql/wal-g-exporter +++ b/puppet/zulip/files/postgresql/wal-g-exporter @@ -3,12 +3,14 @@ import configparser import contextlib import json import logging +import os import subprocess import sys from collections import defaultdict from datetime import datetime, timedelta, timezone from http.server import BaseHTTPRequestHandler, HTTPServer from typing import Dict, List, Mapping, Optional, Protocol +from urllib.parse import parse_qs, urlsplit class GaugeMetric(Protocol): @@ -54,7 +56,7 @@ class WalGPrometheusServer(BaseHTTPRequestHandler): self.wfile.write("\n".join(lines).encode()) def do_GET(self) -> None: - if self.path != "/metrics": + if urlsplit(self.path).path != "/metrics": self.send_response(404) self.end_headers() sys.stderr.flush() @@ -89,13 +91,20 @@ class WalGPrometheusServer(BaseHTTPRequestHandler): now = datetime.now(tz=timezone.utc) try: - config_file = configparser.RawConfigParser() - config_file.read("/etc/zulip/zulip-secrets.conf") - bucket = config_file["secrets"]["s3_backups_bucket"] + query = parse_qs(urlsplit(self.path).query) + modified_env = os.environ.copy() + if query.get("bucket") and len(query["bucket"]) == 1: + bucket = query["bucket"][0] + modified_env["WALG_S3_PREFIX"] = f"s3://{bucket}" + else: + config_file = configparser.RawConfigParser() + config_file.read("/etc/zulip/zulip-secrets.conf") + bucket = config_file["secrets"]["s3_backups_bucket"] backup_list_output = subprocess.check_output( ["env-wal-g", "backup-list", "--detail", "--json"], text=True, + env=modified_env, ) data = json.loads(backup_list_output) backup_count(len(data), {"bucket": bucket})