mikhail
975934373b
list_databases() прячит базы по массиву строк: подстрока без регистра или маска со * ? [. Фильтр — единственное место рождения списка, поэтому одинаков для меню, пульта и list; явное имя в CLI не трогает.
832 lines
33 KiB
Python
832 lines
33 KiB
Python
#!/usr/bin/env python3
|
||
"""Перенос баз PostgreSQL со стенда на локальный сервер: pg_dump -Fc + pg_restore.
|
||
|
||
Два режима: интерактивное меню (запуск без аргументов) и команды CLI.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import fnmatch
|
||
import getpass
|
||
import importlib.util
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
DEFAULT_CONFIG = Path(__file__).with_name("config.json")
|
||
STAMP_RE = re.compile(r"-\d{8}-\d{6}$")
|
||
|
||
# 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
|
||
|
||
|
||
class Cancelled(Exception):
|
||
"""Пользователь отказался от выбора в интерактивном режиме."""
|
||
|
||
|
||
# --------------------------------------------------------------------------- шлюзы наружу
|
||
#
|
||
# Ядро ничего не знает про UI. Три необязательных хука позволяют TUI перехватить то, что
|
||
# иначе ушло бы в stdout или в терминал: строки лога, состояние шага и запрос пароля.
|
||
# Пока хуки не установлены, поведение текстового режима не отличается ни на символ.
|
||
|
||
_sink = None # (kind, text) -> None; kind: log | warn | out
|
||
_progress = None # (event, what, elapsed) -> None; event: start | done | fail
|
||
_password_prompt = None # (label, node) -> str | None; замена getpass под TUI
|
||
|
||
|
||
def set_sink(fn) -> None:
|
||
"""Перенаправить лог и вывод дочерних процессов; None — вернуть печать в stdout."""
|
||
global _sink
|
||
_sink = fn
|
||
|
||
|
||
def set_progress(fn) -> None:
|
||
"""Подписаться на границы шагов (нужно панели плана в TUI)."""
|
||
global _progress
|
||
_progress = fn
|
||
|
||
|
||
def set_password_prompt(fn) -> None:
|
||
"""Заменить getpass: из запущенного Textual читать с терминала нельзя."""
|
||
global _password_prompt
|
||
_password_prompt = fn
|
||
|
||
|
||
def _emit(kind: str, text: str) -> None:
|
||
if _sink is not None:
|
||
_sink(kind, text)
|
||
return
|
||
if kind == "out":
|
||
print(" " + text, flush=True)
|
||
else:
|
||
print(f"[{datetime.now():%H:%M:%S}] {text}", flush=True)
|
||
|
||
|
||
def _notify_progress(event: str, what: str, elapsed: float | None = None) -> None:
|
||
if _progress is not None:
|
||
_progress(event, what, elapsed)
|
||
|
||
|
||
def log(msg: str) -> None:
|
||
_emit("log", msg)
|
||
|
||
|
||
def warn(msg: str) -> None:
|
||
"""То же, что log, но получатель может отличить предупреждение от обычной строки."""
|
||
_emit("warn", msg)
|
||
|
||
|
||
# --------------------------------------------------------------------------- отмена
|
||
#
|
||
# Отмена — это kill дочерних процессов, а не rollback: если DROP DATABASE уже прошёл,
|
||
# локальная база останется пустой. UI обязан говорить об этом прямым текстом.
|
||
|
||
_cancel = threading.Event()
|
||
_procs: set = set()
|
||
_proc_lock = threading.Lock()
|
||
|
||
|
||
def cancel_requested() -> bool:
|
||
return _cancel.is_set()
|
||
|
||
|
||
def clear_cancel() -> None:
|
||
_cancel.clear()
|
||
|
||
|
||
def request_cancel() -> bool:
|
||
"""Прервать текущую операцию. True — если все процессы удалось добить."""
|
||
_cancel.set()
|
||
with _proc_lock:
|
||
victims = list(_procs)
|
||
for proc in victims:
|
||
try:
|
||
proc.terminate()
|
||
except OSError:
|
||
pass
|
||
clean = True
|
||
deadline = time.monotonic() + 3.0
|
||
for proc in victims:
|
||
while proc.poll() is None and time.monotonic() < deadline:
|
||
time.sleep(0.05)
|
||
if proc.poll() is None and not _hard_kill(proc):
|
||
clean = False
|
||
return clean
|
||
|
||
|
||
def _hard_kill(proc) -> bool:
|
||
"""pg_restore -j плодит воркеров: terminate() гасит только родителя, нужно дерево."""
|
||
if os.name == "nt":
|
||
killer = subprocess.run(["taskkill", "/T", "/F", "/PID", str(proc.pid)],
|
||
capture_output=True, text=True, errors="replace")
|
||
if killer.returncode == 0:
|
||
return True
|
||
try:
|
||
proc.kill()
|
||
except OSError:
|
||
pass
|
||
try:
|
||
proc.wait(timeout=2)
|
||
return True
|
||
except subprocess.TimeoutExpired:
|
||
return False
|
||
|
||
|
||
# --------------------------------------------------------------------------- конфиг
|
||
|
||
|
||
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", "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
|
||
|
||
|
||
_password_cache: dict = {}
|
||
|
||
|
||
def _password_key(node: dict, label: str) -> tuple:
|
||
return (label, node.get("host"), node.get("port", 5432), node.get("user"))
|
||
|
||
|
||
def password_needed(node: dict, label: str = "") -> bool:
|
||
"""Придётся ли спрашивать пароль у пользователя (или он уже известен)."""
|
||
if node.get("password"):
|
||
return False
|
||
var = node.get("password_env")
|
||
if var and os.environ.get(var):
|
||
return False
|
||
return _password_key(node, label or node.get("host", "")) not in _password_cache
|
||
|
||
|
||
def prime_password(node: dict, label: str, value: str) -> None:
|
||
"""Положить введённый в UI пароль в общий кэш — второй раз его не спросят."""
|
||
_password_cache[_password_key(node, label)] = value
|
||
|
||
|
||
def forget_password(node: dict, label: str) -> None:
|
||
"""Забыть пароль после отказа сервера, иначе повтор попытки бесполезен."""
|
||
_password_cache.pop(_password_key(node, label), None)
|
||
|
||
|
||
def resolve_password(node: dict, label: str) -> str | None:
|
||
"""Пароль: из конфига, из переменной окружения, иначе спросить один раз за запуск."""
|
||
if node.get("password"):
|
||
return str(node["password"])
|
||
var = node.get("password_env")
|
||
if var:
|
||
value = os.environ.get(var)
|
||
if value:
|
||
return value
|
||
key = _password_key(node, label)
|
||
if key in _password_cache:
|
||
return _password_cache[key]
|
||
if _password_prompt is not None:
|
||
# ленивый вызов может прилететь из середины do_sync — тогда спрашивает UI, не getpass
|
||
value = _password_prompt(label, node)
|
||
if value is None:
|
||
raise Cancelled
|
||
_password_cache[key] = value
|
||
return value
|
||
if not sys.stdin.isatty():
|
||
if var:
|
||
raise ConfigError(f"переменная окружения {var} пуста, а спросить пароль негде")
|
||
return None
|
||
value = getpass.getpass(
|
||
f"пароль {node['user']}@{node['host']} ({label}), Enter — без пароля: ")
|
||
_password_cache[key] = value
|
||
return value
|
||
|
||
|
||
def conn_env(node: dict, label: str = "", dbname: str | None = None) -> dict:
|
||
"""Параметры подключения уходят в окружение, чтобы пароль не светился в командной строке."""
|
||
env = dict(os.environ)
|
||
env["PGHOST"] = str(node["host"])
|
||
env["PGPORT"] = str(node.get("port", 5432))
|
||
env["PGUSER"] = str(node["user"])
|
||
database = dbname or node.get("database") or node.get("maintenance_database") or "postgres"
|
||
env["PGDATABASE"] = str(database)
|
||
password = resolve_password(node, label or node["host"])
|
||
if password:
|
||
env["PGPASSWORD"] = password
|
||
else:
|
||
env.pop("PGPASSWORD", None)
|
||
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 _spawn(cmd: list, env: dict, stderr) -> subprocess.Popen:
|
||
"""Popen вместо run(): вывод нужен построчно по мере появления, а процесс — убиваемым."""
|
||
flags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
|
||
# запуск и регистрация — под одним замком: иначе процесс, стартовавший между снимком
|
||
# victims в request_cancel и добавлением в _procs, переживёт отмену и доработает до конца
|
||
with _proc_lock:
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
proc = subprocess.Popen(cmd, env=env, text=True, encoding="utf-8", errors="replace",
|
||
bufsize=1, stdout=subprocess.PIPE, stderr=stderr,
|
||
creationflags=flags)
|
||
_procs.add(proc)
|
||
return proc
|
||
|
||
|
||
def _forget(proc) -> None:
|
||
with _proc_lock:
|
||
_procs.discard(proc)
|
||
|
||
|
||
# pg_dump новее сервера пишет в дамп SET неизвестных ему GUC (transaction_timeout — с PG 17):
|
||
# restore их не применит и вернёт код 1, хотя данные встали. Строка объясняет красный вывод.
|
||
NEWER_CLIENT_HINT = ("утилиты новее целевого сервера: неизвестные ему SET пропущены — "
|
||
"на данные это не влияет; чтобы убрать шум, снимайте дамп pg_dump той же "
|
||
"мажорной версии, что и локальный сервер")
|
||
|
||
|
||
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
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
_notify_progress("start", what)
|
||
started = time.monotonic()
|
||
proc = _spawn(cmd, env, subprocess.STDOUT)
|
||
explained = False
|
||
try:
|
||
for line in proc.stdout:
|
||
text = line.rstrip("\r\n")
|
||
_emit("out", text)
|
||
if not explained and "unrecognized configuration parameter" in text:
|
||
explained = True
|
||
warn(NEWER_CLIENT_HINT)
|
||
code = proc.wait()
|
||
finally:
|
||
_forget(proc)
|
||
proc.stdout.close()
|
||
elapsed = time.monotonic() - started
|
||
if code != 0:
|
||
_notify_progress("fail", what, elapsed)
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
raise StepError(f"{what} завершился с кодом {code}")
|
||
_notify_progress("done", what, elapsed)
|
||
log(f"{what}: готово за {elapsed:.1f} c")
|
||
|
||
|
||
def no_prompt(env: dict) -> list:
|
||
"""Без пароля в окружении утилиты просят его с терминала и подвисают в скриптах."""
|
||
return [] if env.get("PGPASSWORD") else ["--no-password"]
|
||
|
||
|
||
def capture(cmd: list, env: dict, what: str) -> str:
|
||
# тоже через _spawn: на недоступном стенде psql висит до PGCONNECT_TIMEOUT и должен убиваться
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
proc = _spawn(cmd, env, subprocess.PIPE)
|
||
try:
|
||
out, err = proc.communicate()
|
||
finally:
|
||
_forget(proc)
|
||
if proc.returncode != 0:
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
detail = (err or "").strip().splitlines()
|
||
raise StepError(f"{what}: " + (detail[-1] if detail else f"код {proc.returncode}"))
|
||
return out
|
||
|
||
|
||
def psql_exec(tools: dict, node: dict, dbname: str, sql: str, what: str,
|
||
dry_run: bool, label: str = "") -> None:
|
||
env = conn_env(node, label, dbname)
|
||
cmd = [tools["psql"], "--no-psqlrc", "-v", "ON_ERROR_STOP=1", *no_prompt(env), "-c", sql]
|
||
run(cmd, env, what, dry_run)
|
||
|
||
|
||
def excluded(name: str, patterns) -> bool:
|
||
"""Прячется ли база фильтром: подстрока без регистра, а со звёздочкой — маска."""
|
||
low = name.lower()
|
||
for raw in patterns or []:
|
||
pattern = str(raw).strip().lower()
|
||
if not pattern:
|
||
continue
|
||
if any(ch in pattern for ch in "*?["):
|
||
if fnmatch.fnmatch(low, pattern):
|
||
return True
|
||
elif pattern in low:
|
||
return True
|
||
return False
|
||
|
||
|
||
def list_databases(tools: dict, node: dict, label: str) -> list:
|
||
"""Список баз сервера: (имя, человекочитаемый размер).
|
||
|
||
Отсев по node.exclude_databases делается здесь, в единственном месте, где список
|
||
вообще рождается: и меню, и пульт, и `list` получают его уже без мусорных баз.
|
||
Явное имя базы в CLI фильтр не трогает — оно списком не ходит.
|
||
"""
|
||
admin_db = node.get("maintenance_database") or node.get("database") or "postgres"
|
||
env = conn_env(node, label, admin_db)
|
||
sql = (
|
||
"SELECT datname, "
|
||
" CASE WHEN has_database_privilege(datname, 'CONNECT') "
|
||
" THEN pg_size_pretty(pg_database_size(datname)) ELSE '-' END "
|
||
"FROM pg_database WHERE NOT datistemplate AND datallowconn ORDER BY datname;"
|
||
)
|
||
cmd = [tools["psql"], "--no-psqlrc", "-At", "-F", "\t", *no_prompt(env), "-c", sql]
|
||
out = capture(cmd, env, f"список баз {node['host']}")
|
||
rows = []
|
||
for line in out.splitlines():
|
||
if not line.strip():
|
||
continue
|
||
name, _, size = line.partition("\t")
|
||
if excluded(name, node.get("exclude_databases")):
|
||
continue
|
||
rows.append((name, size or ""))
|
||
return rows
|
||
|
||
|
||
# --------------------------------------------------------------------------- интерактив
|
||
|
||
|
||
def matches(text: str, query: str) -> bool:
|
||
"""Фильтр списков, общий для pick() и TUI: подстроки без регистра, несколько слов — И."""
|
||
if not query:
|
||
return True
|
||
low = text.lower()
|
||
return all(word in low for word in query.lower().split())
|
||
|
||
|
||
def ask(prompt: str) -> str:
|
||
try:
|
||
return input(prompt).strip()
|
||
except EOFError:
|
||
raise Cancelled from None
|
||
|
||
|
||
def pick(items: list, title: str, render, multi: bool = False) -> list:
|
||
"""Список с фильтрацией по вводу: текст сужает список, число(а) выбирают, пустой ввод — сброс."""
|
||
if not items:
|
||
raise StepError(f"{title}: список пуст")
|
||
if not sys.stdin.isatty():
|
||
raise StepError(f"{title}: нет терминала для выбора")
|
||
|
||
flt = ""
|
||
while True:
|
||
shown = [it for it in items if matches(render(it), flt)]
|
||
print()
|
||
print(f"— {title} —" + (f" фильтр: {flt!r}" if flt else ""))
|
||
if not shown:
|
||
print(" ничего не найдено")
|
||
for i, it in enumerate(shown, 1):
|
||
print(f" {i:>3}. {render(it)}")
|
||
hint = "номера через запятую" if multi else "номер"
|
||
choice = ask(f"[{hint} | текст — фильтр | пусто — сброс | q — назад]: ")
|
||
|
||
if choice.lower() in ("q", "выход", "назад"):
|
||
raise Cancelled
|
||
if not choice:
|
||
flt = ""
|
||
continue
|
||
if re.fullmatch(r"[\d\s,]+", choice):
|
||
nums = [int(n) for n in re.split(r"[\s,]+", choice) if n]
|
||
bad = [n for n in nums if not 1 <= n <= len(shown)]
|
||
if bad or not nums:
|
||
print(f" нет такого номера: {', '.join(map(str, bad)) or '-'}")
|
||
continue
|
||
if not multi and len(nums) > 1:
|
||
print(" здесь выбирается один пункт")
|
||
continue
|
||
return [shown[n - 1] for n in nums]
|
||
flt = choice
|
||
|
||
|
||
def confirm(question: str, default: bool = True) -> bool:
|
||
suffix = "[Y/n]" if default else "[y/N]"
|
||
answer = ask(f"{question} {suffix}: ").lower()
|
||
if not answer:
|
||
return default
|
||
return answer in ("y", "yes", "д", "да")
|
||
|
||
|
||
# --------------------------------------------------------------------------- шаги
|
||
|
||
|
||
def dump(cfg: dict, tools: dict, dbname: str, 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)]
|
||
|
||
env = conn_env(src, "стенд", dbname)
|
||
cmd += no_prompt(env)
|
||
out.parent.mkdir(parents=True, exist_ok=True)
|
||
run(cmd, env, f"дамп {dbname}@{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, dbname: str, dry_run: bool) -> None:
|
||
tgt = cfg["target"]
|
||
maintenance = tgt.get("maintenance_database", "postgres")
|
||
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, dbname: str,
|
||
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", dbname]
|
||
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")
|
||
env = conn_env(tgt, "локальная", dbname)
|
||
cmd += no_prompt(env)
|
||
cmd.append(str(dump_file))
|
||
|
||
try:
|
||
run(cmd, env, f"restore в {dbname}@{tgt['host']}", dry_run)
|
||
except StepError as exc:
|
||
# без --exit-on-error pg_restore возвращает 1 на любых игнорируемых ошибках
|
||
if opts.get("exit_on_error", False):
|
||
raise
|
||
warn(f"внимание: {exc} — часть объектов могла не примениться, см. вывод выше")
|
||
|
||
|
||
def post_sql(cfg: dict, tools: dict, dbname: str, dry_run: bool) -> None:
|
||
statements = cfg.get("post_restore_sql") or []
|
||
for i, sql in enumerate(statements, 1):
|
||
psql_exec(tools, cfg["target"], dbname, 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 dump_dir_of(cfg: dict) -> Path:
|
||
dump_dir = Path(cfg.get("dump_dir") or "dumps")
|
||
if not dump_dir.is_absolute():
|
||
dump_dir = Path(__file__).parent / dump_dir
|
||
return dump_dir
|
||
|
||
|
||
def dump_path_for(cfg: dict, dbname: str) -> Path:
|
||
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
return dump_dir_of(cfg) / f"{dbname}-{stamp}.dump"
|
||
|
||
|
||
def db_from_dump_name(path: Path) -> str:
|
||
return STAMP_RE.sub("", path.stem)
|
||
|
||
|
||
def local_dumps(cfg: dict) -> list:
|
||
dump_dir = dump_dir_of(cfg)
|
||
if not dump_dir.exists():
|
||
return []
|
||
return sorted(dump_dir.glob("*.dump"), key=lambda p: p.stat().st_mtime, reverse=True)
|
||
|
||
|
||
def do_dump(cfg: dict, tools: dict, dbnames: list, dry_run: bool) -> list:
|
||
made = []
|
||
for dbname in dbnames:
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
made.append(dump(cfg, tools, dbname, dump_path_for(cfg, dbname), dry_run))
|
||
prune(dump_dir_of(cfg), int(cfg.get("keep_dumps", 5)))
|
||
return made
|
||
|
||
|
||
def do_restore(cfg: dict, tools: dict, dump_file: Path, target_db: str,
|
||
no_recreate: bool, dry_run: bool) -> None:
|
||
recreate = bool(cfg["target"].get("recreate", True)) and not no_recreate
|
||
if recreate:
|
||
recreate_target(cfg, tools, target_db, dry_run)
|
||
restore(cfg, tools, dump_file, target_db, recreate, dry_run)
|
||
post_sql(cfg, tools, target_db, dry_run)
|
||
|
||
|
||
def do_sync(cfg: dict, tools: dict, dbnames: list, target_db: str | None,
|
||
no_recreate: bool, dry_run: bool) -> None:
|
||
for dbname in dbnames:
|
||
if _cancel.is_set():
|
||
raise Cancelled
|
||
dump_file = dump(cfg, tools, dbname, dump_path_for(cfg, dbname), dry_run)
|
||
local = target_db or cfg["target"].get("database") or dbname
|
||
do_restore(cfg, tools, dump_file, local, no_recreate, dry_run)
|
||
prune(dump_dir_of(cfg), int(cfg.get("keep_dumps", 5)))
|
||
|
||
|
||
# --------------------------------------------------------------------------- меню
|
||
|
||
|
||
def choose_source_dbs(cfg: dict, tools: dict, multi: bool) -> list:
|
||
log(f"читаю список баз со стенда {cfg['source']['host']}")
|
||
rows = list_databases(tools, cfg["source"], "стенд")
|
||
chosen = pick(rows, f"базы на {cfg['source']['host']}",
|
||
lambda r: f"{r[0]:<32} {r[1]}", multi=multi)
|
||
return [name for name, _ in chosen]
|
||
|
||
|
||
def choose_dump_file(cfg: dict) -> Path:
|
||
files = local_dumps(cfg)
|
||
if not files:
|
||
raise StepError(f"в {dump_dir_of(cfg)} нет файлов *.dump")
|
||
|
||
def render(p: Path) -> str:
|
||
st = p.stat()
|
||
return (f"{p.name:<44} {st.st_size / 1024 / 1024:>8.1f} МБ "
|
||
f"{datetime.fromtimestamp(st.st_mtime):%d.%m %H:%M}")
|
||
|
||
return pick(files, f"дампы в {dump_dir_of(cfg)}", render)[0]
|
||
|
||
|
||
def ask_target_db(cfg: dict, default: str) -> str:
|
||
answer = ask(f"локальная база [{default}] (Enter — оставить): ")
|
||
return answer or default
|
||
|
||
|
||
def confirm_target(cfg: dict, what: str, target_db: str) -> bool:
|
||
recreate = bool(cfg["target"].get("recreate", True))
|
||
action = (f"пересоздать (DROP+CREATE) локальную {target_db}" if recreate
|
||
else f"накатить поверх локальной {target_db}")
|
||
return confirm(f"{what} → {action}?", default=not recreate)
|
||
|
||
|
||
def interactive(cfg: dict, tools: dict, dry_run: bool) -> int:
|
||
src, tgt = cfg["source"], cfg["target"]
|
||
while True:
|
||
print()
|
||
print("=" * 62)
|
||
print(f" стенд: {src['user']}@{src['host']}:{src.get('port', 5432)}")
|
||
print(f" локально: {tgt['user']}@{tgt['host']}:{tgt.get('port', 5432)}"
|
||
+ (" [dry-run]" if dry_run else ""))
|
||
print("=" * 62)
|
||
print(" 1. Dump + Restore — снять со стенда и залить локально")
|
||
print(" 2. Dump — только снять дамп со стенда")
|
||
print(" 3. Restore — залить локально готовый дамп")
|
||
print(" 0. Выход")
|
||
try:
|
||
choice = ask("выбор: ")
|
||
if choice in ("0", "q", "выход", ""):
|
||
return 0
|
||
|
||
if choice == "1":
|
||
dbs = choose_source_dbs(cfg, tools, multi=False)
|
||
target_db = ask_target_db(cfg, cfg["target"].get("database") or dbs[0])
|
||
if not confirm_target(cfg, dbs[0], target_db):
|
||
continue
|
||
do_sync(cfg, tools, dbs, target_db, no_recreate=False, dry_run=dry_run)
|
||
log("перенос завершён")
|
||
|
||
elif choice == "2":
|
||
dbs = choose_source_dbs(cfg, tools, multi=True)
|
||
do_dump(cfg, tools, dbs, dry_run)
|
||
log("дамп завершён")
|
||
|
||
elif choice == "3":
|
||
dump_file = choose_dump_file(cfg)
|
||
default_db = cfg["target"].get("database") or db_from_dump_name(dump_file)
|
||
target_db = ask_target_db(cfg, default_db)
|
||
if not confirm_target(cfg, dump_file.name, target_db):
|
||
continue
|
||
do_restore(cfg, tools, dump_file, target_db, no_recreate=False, dry_run=dry_run)
|
||
log("восстановление завершено")
|
||
else:
|
||
print("нет такого пункта")
|
||
except Cancelled:
|
||
continue
|
||
except StepError as exc:
|
||
print(f"ОШИБКА: {exc}", file=sys.stderr)
|
||
|
||
|
||
# --------------------------------------------------------------------------- CLI
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
ap = argparse.ArgumentParser(
|
||
description="Перенос БД PostgreSQL со стенда на локальный сервер. "
|
||
"Без команды запускается интерактивное меню.")
|
||
ap.add_argument("-c", "--config", type=Path, default=DEFAULT_CONFIG, help="путь к config.json")
|
||
ap.add_argument("--dry-run", action="store_true", help="показать команды, ничего не выполнять")
|
||
ap.add_argument("--no-tui", action="store_true",
|
||
help="меню текстом, без Textual (для скриптов и отладки)")
|
||
ap.add_argument("--tui", action="store_true",
|
||
help="только TUI: если textual не установлен — ошибка, а не откат в текст")
|
||
sub = ap.add_subparsers(dest="command")
|
||
|
||
p_sync = sub.add_parser("sync", help="dump + restore")
|
||
p_sync.add_argument("db", nargs="*", help="базы на стенде (без имени — выбор из списка)")
|
||
p_sync.add_argument("--target-db", help="имя локальной базы (по умолчанию как на стенде)")
|
||
p_sync.add_argument("--no-recreate", action="store_true",
|
||
help="не пересоздавать локальную базу перед restore")
|
||
|
||
p_dump = sub.add_parser("dump", help="только снять дамп со стенда")
|
||
p_dump.add_argument("db", nargs="*", help="базы на стенде (без имени — выбор из списка)")
|
||
|
||
p_restore = sub.add_parser("restore", help="только залить дамп локально")
|
||
p_restore.add_argument("file", nargs="?", type=Path,
|
||
help="файл .dump (без имени — выбор из списка)")
|
||
p_restore.add_argument("--target-db", help="имя локальной базы")
|
||
p_restore.add_argument("--no-recreate", action="store_true",
|
||
help="не пересоздавать локальную базу перед restore")
|
||
|
||
sub.add_parser("list", help="показать базы на стенде")
|
||
sub.add_parser("dumps", help="показать локальные дампы")
|
||
sub.add_parser("menu", help="интерактивное меню (то же, что запуск без команды)")
|
||
return ap
|
||
|
||
|
||
def tui_available() -> bool:
|
||
return importlib.util.find_spec("textual") is not None
|
||
|
||
|
||
def tui_enabled(args) -> bool:
|
||
"""Единственная точка решения о режиме: условия не размазываются по слоям."""
|
||
if (args.command or "menu") != "menu":
|
||
return False
|
||
if args.no_tui or os.environ.get("PG_STAND_SYNC_NO_TUI"):
|
||
return False
|
||
if not (sys.stdout.isatty() and sys.stdin.isatty()):
|
||
return False
|
||
return tui_available()
|
||
|
||
|
||
def main() -> int:
|
||
args = build_parser().parse_args()
|
||
|
||
try:
|
||
if args.tui and not tui_available():
|
||
raise ConfigError("textual не установлен: pip install -r requirements.txt")
|
||
cfg = load_config(args.config)
|
||
tools = resolve_tools(cfg)
|
||
command = args.command or "menu"
|
||
|
||
if command == "menu":
|
||
if args.tui or tui_enabled(args):
|
||
try:
|
||
import pg_stand_sync_tui as tui
|
||
return tui.run_tui(cfg, tools, args.dry_run)
|
||
except Exception as exc: # TUI не должен лишать работоспособности
|
||
if args.tui:
|
||
raise
|
||
print(f"TUI не поднялся ({exc}), текстовый режим", file=sys.stderr)
|
||
return interactive(cfg, tools, args.dry_run)
|
||
|
||
if command == "list":
|
||
for name, size in list_databases(tools, cfg["source"], "стенд"):
|
||
print(f"{name:<40} {size}")
|
||
return 0
|
||
|
||
if command == "dumps":
|
||
for path in local_dumps(cfg):
|
||
st = path.stat()
|
||
print(f"{path.name:<48} {st.st_size / 1024 / 1024:>8.1f} МБ "
|
||
f"{datetime.fromtimestamp(st.st_mtime):%d.%m.%Y %H:%M}")
|
||
return 0
|
||
|
||
if command in ("sync", "dump"):
|
||
dbs = args.db or choose_source_dbs(cfg, tools, multi=(command == "dump"))
|
||
if command == "dump":
|
||
do_dump(cfg, tools, dbs, args.dry_run)
|
||
log("дамп завершён")
|
||
else:
|
||
do_sync(cfg, tools, dbs, args.target_db, args.no_recreate, args.dry_run)
|
||
log("перенос завершён")
|
||
return 0
|
||
|
||
if command == "restore":
|
||
dump_file = args.file or choose_dump_file(cfg)
|
||
target_db = (args.target_db or cfg["target"].get("database")
|
||
or db_from_dump_name(dump_file))
|
||
do_restore(cfg, tools, dump_file, target_db, args.no_recreate, args.dry_run)
|
||
log("восстановление завершено")
|
||
return 0
|
||
|
||
return 0
|
||
except (ConfigError, StepError) as exc:
|
||
print(f"ОШИБКА: {exc}", file=sys.stderr)
|
||
return 1
|
||
except Cancelled:
|
||
print("отменено", file=sys.stderr)
|
||
return 130
|
||
except KeyboardInterrupt:
|
||
print("прервано пользователем", file=sys.stderr)
|
||
return 130
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|