mikhail
267 lines
10 KiB
Python
267 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Перенос базы PostgreSQL со стенда на локальный сервер: pg_dump -Fc + pg_restore."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
DEFAULT_CONFIG = Path(__file__).with_name("config.json")
|
|
|
|
# Windows-консоль по умолчанию отдаёт cp866, русский лог в ней превращается в мусор
|
|
for stream in (sys.stdout, sys.stderr):
|
|
if hasattr(stream, "reconfigure"):
|
|
stream.reconfigure(encoding="utf-8", errors="replace")
|
|
|
|
|
|
class ConfigError(Exception):
|
|
pass
|
|
|
|
|
|
class StepError(Exception):
|
|
pass
|
|
|
|
|
|
def log(msg: str) -> None:
|
|
print(f"[{datetime.now():%H:%M:%S}] {msg}", flush=True)
|
|
|
|
|
|
def load_config(path: Path) -> dict:
|
|
if not path.exists():
|
|
raise ConfigError(
|
|
f"конфиг не найден: {path}\n"
|
|
f"скопируйте config.example.json в {path.name} и заполните"
|
|
)
|
|
try:
|
|
cfg = json.loads(path.read_text(encoding="utf-8-sig"))
|
|
except json.JSONDecodeError as exc:
|
|
raise ConfigError(f"{path}: не разбирается как JSON — {exc}") from exc
|
|
|
|
for section in ("source", "target"):
|
|
node = cfg.get(section)
|
|
if not isinstance(node, dict):
|
|
raise ConfigError(f'{path}: нет секции "{section}"')
|
|
for key in ("host", "database", "user"):
|
|
if not node.get(key):
|
|
raise ConfigError(f"{path}: {section}.{key} не задан")
|
|
return cfg
|
|
|
|
|
|
def resolve_tools(cfg: dict) -> dict:
|
|
bin_dir = cfg.get("pg_bin_dir") or ""
|
|
tools = {}
|
|
for name in ("pg_dump", "pg_restore", "psql"):
|
|
if bin_dir:
|
|
candidate = Path(bin_dir) / (name + (".exe" if os.name == "nt" else ""))
|
|
if not candidate.exists():
|
|
raise ConfigError(f"не найден {name} в pg_bin_dir: {bin_dir}")
|
|
tools[name] = str(candidate)
|
|
else:
|
|
found = shutil.which(name)
|
|
if not found:
|
|
raise ConfigError(f"{name} не найден в PATH, задайте pg_bin_dir в конфиге")
|
|
tools[name] = found
|
|
return tools
|
|
|
|
|
|
def conn_env(node: dict) -> dict:
|
|
"""Параметры подключения уходят в окружение, чтобы пароль не светился в командной строке."""
|
|
env = dict(os.environ)
|
|
env["PGHOST"] = str(node["host"])
|
|
env["PGPORT"] = str(node.get("port", 5432))
|
|
env["PGUSER"] = str(node["user"])
|
|
env["PGDATABASE"] = str(node["database"])
|
|
password = node.get("password")
|
|
if password:
|
|
env["PGPASSWORD"] = str(password)
|
|
elif node.get("password_env"):
|
|
value = os.environ.get(node["password_env"])
|
|
if not value:
|
|
raise ConfigError(f"переменная окружения {node['password_env']} пуста")
|
|
env["PGPASSWORD"] = value
|
|
if node.get("sslmode"):
|
|
env["PGSSLMODE"] = str(node["sslmode"])
|
|
env.setdefault("PGCLIENTENCODING", "UTF8")
|
|
env["PGCONNECT_TIMEOUT"] = str(node.get("connect_timeout", 15))
|
|
return env
|
|
|
|
|
|
def run(cmd: list, env: dict, what: str, dry_run: bool) -> None:
|
|
shown = " ".join(f'"{c}"' if " " in c else c for c in cmd)
|
|
log(f"{what}: {shown}")
|
|
if dry_run:
|
|
return
|
|
started = time.monotonic()
|
|
proc = subprocess.run(cmd, env=env, text=True, encoding="utf-8", errors="replace",
|
|
stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
if proc.stdout:
|
|
for line in proc.stdout.splitlines():
|
|
print(" " + line, flush=True)
|
|
if proc.returncode != 0:
|
|
raise StepError(f"{what} завершился с кодом {proc.returncode}")
|
|
log(f"{what}: готово за {time.monotonic() - started:.1f} c")
|
|
|
|
|
|
def psql_exec(tools: dict, node: dict, dbname: str, sql: str, what: str, dry_run: bool) -> None:
|
|
env = conn_env({**node, "database": dbname})
|
|
cmd = [tools["psql"], "--no-psqlrc", "-v", "ON_ERROR_STOP=1", "-c", sql]
|
|
run(cmd, env, what, dry_run)
|
|
|
|
|
|
def dump(cfg: dict, tools: dict, out: Path, dry_run: bool) -> Path:
|
|
src = cfg["source"]
|
|
opts = cfg.get("dump", {})
|
|
cmd = [tools["pg_dump"], "--format=custom"]
|
|
if opts.get("verbose"):
|
|
cmd.append("--verbose")
|
|
if opts.get("compress") is not None:
|
|
cmd.append(f"--compress={opts['compress']}")
|
|
if opts.get("schema_only"):
|
|
cmd.append("--schema-only")
|
|
if opts.get("no_owner", True):
|
|
cmd.append("--no-owner")
|
|
if opts.get("no_privileges", True):
|
|
cmd.append("--no-privileges")
|
|
for schema in opts.get("schemas", []):
|
|
cmd += ["--schema", schema]
|
|
for schema in opts.get("exclude_schemas", []):
|
|
cmd += ["--exclude-schema", schema]
|
|
for table in opts.get("exclude_tables", []):
|
|
cmd += ["--exclude-table", table]
|
|
for table in opts.get("exclude_table_data", []):
|
|
cmd += ["--exclude-table-data", table]
|
|
cmd += ["--file", str(out)]
|
|
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
run(cmd, conn_env(src), f"дамп {src['database']}@{src['host']}", dry_run)
|
|
if not dry_run:
|
|
size = out.stat().st_size / 1024 / 1024
|
|
log(f"дамп: {out} ({size:.1f} МБ)")
|
|
return out
|
|
|
|
|
|
def recreate_target(cfg: dict, tools: dict, dry_run: bool) -> None:
|
|
tgt = cfg["target"]
|
|
maintenance = tgt.get("maintenance_database", "postgres")
|
|
dbname = tgt["database"]
|
|
quoted = '"' + dbname.replace('"', '""') + '"'
|
|
literal = "'" + dbname.replace("'", "''") + "'"
|
|
|
|
log(f"пересоздаю локальную базу {dbname}")
|
|
psql_exec(tools, tgt, maintenance,
|
|
f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
|
|
f"WHERE datname = {literal} AND pid <> pg_backend_pid();",
|
|
"обрыв активных сессий", dry_run)
|
|
psql_exec(tools, tgt, maintenance, f"DROP DATABASE IF EXISTS {quoted};",
|
|
"DROP DATABASE", dry_run)
|
|
create = f"CREATE DATABASE {quoted}"
|
|
if tgt.get("owner"):
|
|
create += ' OWNER "' + tgt["owner"].replace('"', '""') + '"'
|
|
if tgt.get("template"):
|
|
create += f" TEMPLATE {tgt['template']}"
|
|
psql_exec(tools, tgt, maintenance, create + ";", "CREATE DATABASE", dry_run)
|
|
|
|
|
|
def restore(cfg: dict, tools: dict, dump_file: Path, recreated: bool, dry_run: bool) -> None:
|
|
tgt = cfg["target"]
|
|
opts = cfg.get("restore", {})
|
|
if not dry_run and not dump_file.exists():
|
|
raise StepError(f"файл дампа не найден: {dump_file}")
|
|
|
|
cmd = [tools["pg_restore"], "--dbname", tgt["database"]]
|
|
jobs = int(opts.get("jobs", 4))
|
|
if jobs > 1:
|
|
cmd += ["--jobs", str(jobs)]
|
|
if opts.get("no_owner", True):
|
|
cmd.append("--no-owner")
|
|
if opts.get("no_privileges", True):
|
|
cmd.append("--no-privileges")
|
|
if opts.get("clean", True) and not recreated:
|
|
cmd += ["--clean", "--if-exists"]
|
|
if opts.get("exit_on_error", False):
|
|
cmd.append("--exit-on-error")
|
|
if opts.get("verbose"):
|
|
cmd.append("--verbose")
|
|
cmd.append(str(dump_file))
|
|
|
|
try:
|
|
run(cmd, conn_env(tgt), f"restore в {tgt['database']}@{tgt['host']}", dry_run)
|
|
except StepError as exc:
|
|
# без --exit-on-error pg_restore возвращает 1 на любых игнорируемых ошибках
|
|
if opts.get("exit_on_error", False):
|
|
raise
|
|
log(f"внимание: {exc} — часть объектов могла не примениться, см. вывод выше")
|
|
|
|
|
|
def post_sql(cfg: dict, tools: dict, dry_run: bool) -> None:
|
|
statements = cfg.get("post_restore_sql") or []
|
|
for i, sql in enumerate(statements, 1):
|
|
psql_exec(tools, cfg["target"], cfg["target"]["database"], sql,
|
|
f"post-SQL {i}/{len(statements)}", dry_run)
|
|
|
|
|
|
def prune(dump_dir: Path, keep: int) -> None:
|
|
if keep <= 0 or not dump_dir.exists():
|
|
return
|
|
files = sorted(dump_dir.glob("*.dump"), key=lambda p: p.stat().st_mtime, reverse=True)
|
|
for old in files[keep:]:
|
|
log(f"удаляю старый дамп {old.name}")
|
|
old.unlink()
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description="Перенос БД PostgreSQL со стенда на локальный сервер")
|
|
ap.add_argument("-c", "--config", type=Path, default=DEFAULT_CONFIG, help="путь к config.json")
|
|
ap.add_argument("--dump-only", action="store_true", help="только снять дамп со стенда")
|
|
ap.add_argument("--restore-only", metavar="FILE", type=Path,
|
|
help="только восстановить указанный дамп локально")
|
|
ap.add_argument("--no-recreate", action="store_true",
|
|
help="не пересоздавать локальную базу перед restore")
|
|
ap.add_argument("--dry-run", action="store_true", help="показать команды, ничего не выполнять")
|
|
args = ap.parse_args()
|
|
|
|
try:
|
|
cfg = load_config(args.config)
|
|
tools = resolve_tools(cfg)
|
|
dump_dir = Path(cfg.get("dump_dir") or "dumps")
|
|
if not dump_dir.is_absolute():
|
|
dump_dir = Path(__file__).parent / dump_dir
|
|
|
|
if args.restore_only:
|
|
dump_file = args.restore_only
|
|
else:
|
|
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
dump_file = dump_dir / f"{cfg['source']['database']}-{stamp}.dump"
|
|
dump(cfg, tools, dump_file, args.dry_run)
|
|
if args.dump_only:
|
|
log("режим --dump-only: restore пропущен")
|
|
prune(dump_dir, int(cfg.get("keep_dumps", 5)))
|
|
return 0
|
|
|
|
recreate = bool(cfg["target"].get("recreate", True)) and not args.no_recreate
|
|
if recreate:
|
|
recreate_target(cfg, tools, args.dry_run)
|
|
restore(cfg, tools, dump_file, recreate, args.dry_run)
|
|
post_sql(cfg, tools, args.dry_run)
|
|
if not args.restore_only:
|
|
prune(dump_dir, int(cfg.get("keep_dumps", 5)))
|
|
log("перенос завершён")
|
|
return 0
|
|
except (ConfigError, StepError) as exc:
|
|
print(f"ОШИБКА: {exc}", file=sys.stderr)
|
|
return 1
|
|
except KeyboardInterrupt:
|
|
print("прервано пользователем", file=sys.stderr)
|
|
return 130
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|