445 lines
17 KiB
Python
445 lines
17 KiB
Python
"""Asynchronous PostgreSQL backup and restore jobs for Inventera."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from copy import deepcopy
|
||
from datetime import datetime, timedelta
|
||
import glob
|
||
import os
|
||
from pathlib import Path
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import uuid
|
||
|
||
|
||
class DatabaseMaintenanceError(RuntimeError):
|
||
"""Base error exposed to the database maintenance API."""
|
||
|
||
|
||
class DatabaseMaintenanceBusy(DatabaseMaintenanceError):
|
||
"""Raised when another backup or restore job is already running."""
|
||
|
||
|
||
class PostgreSqlToolNotFound(DatabaseMaintenanceError):
|
||
"""Raised when a required PostgreSQL command-line tool is unavailable."""
|
||
|
||
|
||
class InvalidBackupFile(DatabaseMaintenanceError):
|
||
"""Raised when the uploaded file is not a PostgreSQL custom-format dump."""
|
||
|
||
|
||
class DatabaseMaintenanceManager:
|
||
def __init__(self, db_config, work_directory=None):
|
||
self._db_config = dict(db_config)
|
||
self._work_directory = Path(work_directory or self._default_work_directory())
|
||
self._jobs = {}
|
||
self._active_job_id = None
|
||
self._lock = threading.RLock()
|
||
|
||
@staticmethod
|
||
def _default_work_directory():
|
||
program_data = os.environ.get("PROGRAMDATA")
|
||
if program_data:
|
||
return os.path.join(program_data, "Inventera", "database-maintenance")
|
||
return os.path.join(tempfile.gettempdir(), "inventera", "database-maintenance")
|
||
|
||
@staticmethod
|
||
def _find_postgresql_tool(tool_name):
|
||
executable_name = f"{tool_name}.exe" if os.name == "nt" else tool_name
|
||
discovered = shutil.which(executable_name) or shutil.which(tool_name)
|
||
if discovered:
|
||
return discovered
|
||
|
||
candidates = []
|
||
|
||
discovered_psql = shutil.which("psql.exe" if os.name == "nt" else "psql")
|
||
if discovered_psql:
|
||
sibling = os.path.join(os.path.dirname(discovered_psql), executable_name)
|
||
if os.path.isfile(sibling):
|
||
candidates.append(sibling)
|
||
|
||
windows_roots = []
|
||
for environment_name in ("ProgramW6432", "ProgramFiles", "ProgramFiles(x86)"):
|
||
value = os.environ.get(environment_name)
|
||
if value and value not in windows_roots:
|
||
windows_roots.append(value)
|
||
|
||
for root in windows_roots:
|
||
candidates.extend(
|
||
glob.glob(os.path.join(root, "PostgreSQL", "*", "bin", executable_name))
|
||
)
|
||
|
||
common_patterns = [
|
||
os.path.join("/Library/PostgreSQL", "*", "bin", executable_name),
|
||
os.path.join(
|
||
"/Applications/Postgres.app/Contents/Versions",
|
||
"*",
|
||
"bin",
|
||
executable_name,
|
||
),
|
||
os.path.join("/opt/homebrew/opt", "postgresql*", "bin", executable_name),
|
||
os.path.join("/usr/local/opt", "postgresql*", "bin", executable_name),
|
||
os.path.join("/usr/lib/postgresql", "*", "bin", executable_name),
|
||
]
|
||
for pattern in common_patterns:
|
||
candidates.extend(glob.glob(pattern))
|
||
|
||
def version_key(candidate):
|
||
version_name = Path(candidate).parent.parent.name
|
||
return tuple(int(part) for part in re.findall(r"\d+", version_name))
|
||
|
||
def candidate_key(candidate):
|
||
normalized = candidate.replace("\\", "/").lower()
|
||
if "/library/postgresql/" in normalized:
|
||
distribution_priority = 4
|
||
elif "/applications/postgres.app/" in normalized:
|
||
distribution_priority = 3
|
||
elif "/postgresql/" in normalized:
|
||
distribution_priority = 2
|
||
else:
|
||
distribution_priority = 1
|
||
return version_key(candidate), distribution_priority, normalized
|
||
|
||
candidates = sorted(set(candidates), key=candidate_key, reverse=True)
|
||
for candidate in candidates:
|
||
if os.path.isfile(candidate):
|
||
return candidate
|
||
|
||
raise PostgreSqlToolNotFound(
|
||
f"Не найден {executable_name}. Установите командные инструменты PostgreSQL."
|
||
)
|
||
|
||
@staticmethod
|
||
def _popen_options():
|
||
options = {}
|
||
if os.name == "nt":
|
||
options["creationflags"] = subprocess.CREATE_NO_WINDOW
|
||
return options
|
||
|
||
def _postgres_environment(self):
|
||
environment = os.environ.copy()
|
||
environment["PGPASSWORD"] = str(self._db_config.get("password", ""))
|
||
environment["PGCLIENTENCODING"] = "UTF8"
|
||
return environment
|
||
|
||
def _connection_arguments(self):
|
||
return [
|
||
"--host",
|
||
str(self._db_config.get("host", "localhost")),
|
||
"--port",
|
||
str(self._db_config.get("port", 5432)),
|
||
"--username",
|
||
str(self._db_config.get("user", "postgres")),
|
||
"--dbname",
|
||
str(self._db_config.get("dbname", "postgres")),
|
||
"--no-password",
|
||
]
|
||
|
||
def _ensure_work_directories(self):
|
||
backup_directory = self._work_directory / "backups"
|
||
upload_directory = self._work_directory / "uploads"
|
||
backup_directory.mkdir(parents=True, exist_ok=True)
|
||
upload_directory.mkdir(parents=True, exist_ok=True)
|
||
return backup_directory, upload_directory
|
||
|
||
def _prune_jobs_locked(self):
|
||
cutoff = datetime.now() - timedelta(hours=24)
|
||
expired_ids = [
|
||
job_id
|
||
for job_id, job in self._jobs.items()
|
||
if job.get("created_at") and job["created_at"] < cutoff
|
||
and job.get("status") not in ("queued", "running")
|
||
]
|
||
for job_id in expired_ids:
|
||
job = self._jobs.pop(job_id)
|
||
path = job.get("path")
|
||
if path:
|
||
try:
|
||
Path(path).unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|
||
|
||
def _create_job(self, kind, owner, filename=None, path=None):
|
||
with self._lock:
|
||
self._prune_jobs_locked()
|
||
if self._active_job_id:
|
||
active = self._jobs.get(self._active_job_id)
|
||
if active and active.get("status") in ("queued", "running"):
|
||
raise DatabaseMaintenanceBusy(
|
||
"Дождитесь завершения текущей операции с базой данных."
|
||
)
|
||
self._active_job_id = None
|
||
|
||
job_id = uuid.uuid4().hex
|
||
self._jobs[job_id] = {
|
||
"id": job_id,
|
||
"kind": kind,
|
||
"owner": owner,
|
||
"status": "queued",
|
||
"progress": 0,
|
||
"message": "Операция поставлена в очередь.",
|
||
"filename": filename,
|
||
"path": str(path) if path else None,
|
||
"created_at": datetime.now(),
|
||
}
|
||
self._active_job_id = job_id
|
||
return job_id
|
||
|
||
def _update_job(self, job_id, **updates):
|
||
with self._lock:
|
||
job = self._jobs.get(job_id)
|
||
if not job:
|
||
return
|
||
job.update(updates)
|
||
if job.get("status") not in ("queued", "running") and self._active_job_id == job_id:
|
||
self._active_job_id = None
|
||
|
||
def _fail_job(self, job_id, message):
|
||
self._update_job(
|
||
job_id,
|
||
status="error",
|
||
progress=100,
|
||
message=message or "Операция завершилась с ошибкой.",
|
||
)
|
||
|
||
def get_job(self, job_id, owner):
|
||
with self._lock:
|
||
job = self._jobs.get(job_id)
|
||
if not job or job.get("owner") != owner:
|
||
return None
|
||
public_job = deepcopy(job)
|
||
public_job.pop("owner", None)
|
||
public_job.pop("path", None)
|
||
created_at = public_job.get("created_at")
|
||
if created_at:
|
||
public_job["created_at"] = created_at.isoformat(timespec="seconds")
|
||
public_job["download_ready"] = (
|
||
public_job.get("kind") == "backup" and public_job.get("status") == "success"
|
||
)
|
||
return public_job
|
||
|
||
def get_backup_download(self, job_id, owner):
|
||
with self._lock:
|
||
job = self._jobs.get(job_id)
|
||
if (
|
||
not job
|
||
or job.get("owner") != owner
|
||
or job.get("kind") != "backup"
|
||
or job.get("status") != "success"
|
||
):
|
||
return None
|
||
path = job.get("path")
|
||
filename = job.get("filename")
|
||
if not path or not os.path.isfile(path):
|
||
return None
|
||
return path, filename
|
||
|
||
def start_backup(self, owner, estimated_database_size=0):
|
||
pg_dump = self._find_postgresql_tool("pg_dump")
|
||
backup_directory, _ = self._ensure_work_directories()
|
||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
filename = f"Inventera_{timestamp}.backup"
|
||
backup_path = backup_directory / f"{uuid.uuid4().hex}_{filename}"
|
||
job_id = self._create_job("backup", owner, filename, backup_path)
|
||
worker = threading.Thread(
|
||
target=self._run_backup,
|
||
args=(job_id, pg_dump, backup_path, max(int(estimated_database_size or 0), 0)),
|
||
name=f"inventera-backup-{job_id[:8]}",
|
||
daemon=True,
|
||
)
|
||
worker.start()
|
||
return job_id
|
||
|
||
def _run_backup(self, job_id, pg_dump, backup_path, estimated_database_size):
|
||
error_path = backup_path.with_suffix(".log")
|
||
try:
|
||
self._update_job(
|
||
job_id,
|
||
status="running",
|
||
progress=3,
|
||
message="Создание резервной копии базы данных...",
|
||
)
|
||
command = [
|
||
pg_dump,
|
||
*self._connection_arguments(),
|
||
"--format=custom",
|
||
"--no-owner",
|
||
"--no-privileges",
|
||
"--file",
|
||
str(backup_path),
|
||
]
|
||
with open(error_path, "wb") as error_file:
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=error_file,
|
||
env=self._postgres_environment(),
|
||
**self._popen_options(),
|
||
)
|
||
fallback_progress = 3
|
||
while process.poll() is None:
|
||
if estimated_database_size > 0 and backup_path.exists():
|
||
generated_size = backup_path.stat().st_size
|
||
progress = min(92, 5 + int((generated_size / estimated_database_size) * 87))
|
||
else:
|
||
fallback_progress = min(88, fallback_progress + 1)
|
||
progress = fallback_progress
|
||
self._update_job(job_id, progress=progress)
|
||
time.sleep(0.5)
|
||
exit_code = process.returncode
|
||
|
||
if exit_code != 0:
|
||
details = self._read_error_details(error_path)
|
||
raise DatabaseMaintenanceError(
|
||
details or f"pg_dump завершился с кодом {exit_code}."
|
||
)
|
||
if not backup_path.exists() or backup_path.stat().st_size <= 5:
|
||
raise DatabaseMaintenanceError("PostgreSQL не создал файл резервной копии.")
|
||
|
||
self._update_job(
|
||
job_id,
|
||
status="success",
|
||
progress=100,
|
||
message="Резервная копия создана. Скачивание файла начнётся автоматически.",
|
||
)
|
||
except Exception as exc:
|
||
try:
|
||
backup_path.unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|
||
self._fail_job(job_id, f"Ошибка резервного копирования: {exc}")
|
||
finally:
|
||
try:
|
||
error_path.unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|
||
|
||
@staticmethod
|
||
def _read_error_details(error_path):
|
||
try:
|
||
text = error_path.read_text(encoding="utf-8", errors="replace").strip()
|
||
except OSError:
|
||
return ""
|
||
if len(text) > 1600:
|
||
text = text[-1600:]
|
||
return text
|
||
|
||
def start_restore(self, owner, uploaded_file, original_filename):
|
||
pg_restore = self._find_postgresql_tool("pg_restore")
|
||
_, upload_directory = self._ensure_work_directories()
|
||
suffix = Path(original_filename or "").suffix.lower()
|
||
if suffix not in (".backup", ".dump"):
|
||
raise InvalidBackupFile(
|
||
"Выберите резервную копию PostgreSQL с расширением .backup или .dump."
|
||
)
|
||
|
||
upload_path = upload_directory / f"{uuid.uuid4().hex}{suffix}"
|
||
job_id = self._create_job("restore", owner, original_filename or upload_path.name, upload_path)
|
||
try:
|
||
uploaded_file.save(upload_path)
|
||
with open(upload_path, "rb") as backup_file:
|
||
signature = backup_file.read(5)
|
||
if signature != b"PGDMP":
|
||
raise InvalidBackupFile(
|
||
"Файл не является резервной копией PostgreSQL, созданной Inventera."
|
||
)
|
||
except Exception:
|
||
try:
|
||
upload_path.unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|
||
self._update_job(job_id, status="error", progress=100)
|
||
raise
|
||
|
||
worker = threading.Thread(
|
||
target=self._run_restore,
|
||
args=(job_id, pg_restore, upload_path),
|
||
name=f"inventera-restore-{job_id[:8]}",
|
||
daemon=True,
|
||
)
|
||
worker.start()
|
||
return job_id
|
||
|
||
def _restore_item_count(self, pg_restore, upload_path):
|
||
result = subprocess.run(
|
||
[pg_restore, "--list", str(upload_path)],
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.DEVNULL,
|
||
env=self._postgres_environment(),
|
||
check=False,
|
||
**self._popen_options(),
|
||
)
|
||
if result.returncode != 0:
|
||
return 1
|
||
text = result.stdout.decode("utf-8", errors="replace")
|
||
return max(1, sum(1 for line in text.splitlines() if line and not line.startswith(";")))
|
||
|
||
def _run_restore(self, job_id, pg_restore, upload_path):
|
||
error_lines = []
|
||
try:
|
||
self._update_job(
|
||
job_id,
|
||
status="running",
|
||
progress=3,
|
||
message="Восстановление базы данных...",
|
||
)
|
||
total_items = self._restore_item_count(pg_restore, upload_path)
|
||
command = [
|
||
pg_restore,
|
||
*self._connection_arguments(),
|
||
"--clean",
|
||
"--if-exists",
|
||
"--no-owner",
|
||
"--no-privileges",
|
||
"--exit-on-error",
|
||
"--single-transaction",
|
||
"--verbose",
|
||
str(upload_path),
|
||
]
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.PIPE,
|
||
env=self._postgres_environment(),
|
||
**self._popen_options(),
|
||
)
|
||
processed_items = 0
|
||
if process.stderr is not None:
|
||
for raw_line in iter(process.stderr.readline, b""):
|
||
line = raw_line.decode("utf-8", errors="replace").strip()
|
||
if line:
|
||
error_lines.append(line)
|
||
if len(error_lines) > 30:
|
||
error_lines.pop(0)
|
||
processed_items += 1
|
||
progress = min(96, 5 + int((processed_items / total_items) * 90))
|
||
self._update_job(
|
||
job_id,
|
||
progress=progress,
|
||
message="Восстановление объектов базы данных...",
|
||
)
|
||
exit_code = process.wait()
|
||
if exit_code != 0:
|
||
details = "\n".join(error_lines[-12:])
|
||
raise DatabaseMaintenanceError(
|
||
details or f"pg_restore завершился с кодом {exit_code}."
|
||
)
|
||
|
||
self._update_job(
|
||
job_id,
|
||
status="success",
|
||
progress=100,
|
||
message="Восстановление завершено успешно.",
|
||
)
|
||
except Exception as exc:
|
||
self._fail_job(job_id, f"Ошибка восстановления базы данных: {exc}")
|
||
finally:
|
||
try:
|
||
upload_path.unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|