import os.path
import pathlib
from subprocess import PIPE, DEVNULL
import orjson
from cattrs import unstructure
from ferrosoft.ops import commands, sql
from ferrosoft.ops.models import (
Journal,
JournalEntry,
Tenant,
get_tenants,
get_roles,
get_session_user,
get_databases,
)
[docs]
class LocalFileError(RuntimeError):
pass
def _get_journal_contents(
journal_file_name: str,
aws_executable: str = "aws",
bucket_name: str = "backup",
local_mode: bool = False,
) -> str:
if local_mode:
with open(journal_file_name, "r") as journal_file:
return journal_file.read()
else:
result = commands.run(
[
aws_executable,
"s3",
"cp",
"s3://%s/%s" % (bucket_name, journal_file_name),
"-",
],
stdout=PIPE,
encoding="utf-8",
errors="replace",
)
return result.stdout
[docs]
def get_journal(*args, **kwargs) -> Journal:
content = _get_journal_contents(*args, **kwargs)
return unstructure(orjson.loads(content), Journal)
[docs]
def fetch_backups(
journal: Journal,
local_dir: str,
aws_executable: str = "aws",
bucket_name: str = "backup",
local_mode: bool = False,
):
"""
Retrieve backup files from object storage and save them to a local directory.
"""
for entry in journal.entries:
dest_filename = os.path.join(local_dir, entry.filename)
if local_mode:
local_file = pathlib.Path(dest_filename)
if not local_file.exists:
raise LocalFileError("File not found: %s" % dest_filename)
else:
with open(dest_filename, mode="wb") as dest_file:
commands.run(
[
aws_executable,
"s3",
"cp",
"s3://%s/%s" % (bucket_name, entry.filename),
"-",
],
stdout=dest_file,
)
[docs]
def decompress_backups(
journal: Journal,
local_dir: str,
decompress_executable: str = "unzstd",
):
"""
Decompress backup files found in local directory.
Args:
journal:
local_dir:
decompress_executable: A decompression executable file name which must support -c option.
Returns:
"""
for entry in journal.entries:
if decompress_executable == "":
# No decompression needed
entry.filename_decompressed = entry.filename
continue
# When a file is compressed, assume the last file extension
# indicates the compression algorithm. It must be stripped.
file_path = pathlib.Path(os.path.join(local_dir, entry.filename))
entry.filename_decompressed = os.path.join(local_dir, file_path.stem)
with open(entry.filename_decompressed, mode="wb") as dest_file:
commands.run(
[
decompress_executable,
"-c",
str(file_path),
],
stdout=dest_file,
)
[docs]
def restore_databases(
journal: Journal,
psql_executable: str,
local_mode: bool = False,
master_db: str | None = None,
):
# Restore master database before tenant databases because master database
# contains tenant database passwords.
master_entry = journal.master_entry()
if master_entry is not None:
restore_database_entry(master_entry, psql_executable, database_name=master_db)
# Load tenants from freshly restored master database, because it contains
# tenant passwords.
for tenant in get_tenants(psql_executable):
tenant_entry = journal.tenant_entry(tenant.id)
tenant_entry.tenant.db_password = tenant.db_password
restore_database_entry(tenant_entry, psql_executable, local_mode)
def _can_manage_tenant_resources(local_mode: bool, tenant_db_name: str):
# If PGHOST equals the host name of the tenant database, assume
# current session postgres user can create the tenant role and database.
return local_mode or tenant_db_name == os.getenv("PGHOST", "")
[docs]
def restore_database_entry(
entry: JournalEntry,
psql_executable: str,
local_mode: bool = False,
database_name: str | None = None,
):
if entry.filename_decompressed is None or entry.filename_decompressed == "":
raise ValueError("Decompressed filename is required")
psql_args = []
psql_env = None
if entry.tenant is not None:
if _can_manage_tenant_resources(local_mode, entry.tenant.db_hostname):
create_tenant_role(entry.tenant, psql_executable)
create_tenant_database(entry.tenant, psql_executable)
if not local_mode:
psql_args.extend(
[
"--host",
entry.tenant.db_hostname,
"--port",
entry.tenant.db_port,
"--username",
entry.tenant.db_username,
]
)
psql_env = {"PGPASSWORD": entry.tenant.db_password}
with open(entry.filename_decompressed, mode="rb") as backup_file:
commands.run(
[psql_executable, *psql_args, database_name or entry.db_name],
env=psql_env,
stdin=backup_file,
stdout=DEVNULL,
)
if entry.tenant is not None and _can_manage_tenant_resources(
local_mode, entry.tenant.db_hostname
):
assign_tenant_db_owner(entry.tenant, psql_executable)
[docs]
def create_tenant_role(tenant: Tenant, psql_executable: str):
if tenant.db_password == "":
raise ValueError("tenant database password is required")
roles = get_roles(psql_executable)
if tenant.db_username in roles:
return
session_user = get_session_user(psql_executable)
commands.run(
[
psql_executable,
"-c",
sql.create_tenant_role_command(
tenant.db_username, tenant.db_password, session_user
),
]
)
[docs]
def create_tenant_database(tenant: Tenant, psql_executable: str):
databases = get_databases(psql_executable)
if tenant.db_name in databases:
return
commands.run(
[
psql_executable,
"-c",
sql.create_tenant_database_command(tenant.db_name, tenant.db_username),
]
)
[docs]
def assign_tenant_db_owner(tenant: Tenant, psql_executable: str):
commands.run(
[
psql_executable,
"-c",
sql.reassign_owner(tenant.db_username),
tenant.db_name,
]
)
[docs]
def set_tenant_localhost(psql_executable: str, postgres_localhost: str):
commands.run(
[
psql_executable,
"-c",
sql.tenant_host_update_command(postgres_localhost),
]
)