1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
|
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Utilities to help with PostgreSQL databases.
Uses utilities from 'postgres' package such as 'psql' and 'pg_dump'.
"""
import os
import pathlib
from plinth import action_utils
def _run_as(command, **kwargs):
"""Run a command as 'postgres' user."""
command = ['sudo', '--user', 'postgres'] + command
return action_utils.run(command, check=True, **kwargs)
def run_query(query):
"""Run a database query as 'postgres' user.
Does not ensure that database server is running.
"""
env = os.environ.copy()
env['ON_ERROR_EXIT'] = '1'
return _run_as(['psql', '--echo-errors'], env=env,
input=query.encode('utf-8'))
def _create_user(database_user: str, database_password: str):
"""Create a new user account with given credentials.
Ignore errors if user already exists. Set password on the account either
way. Passwords must be alphanumeric.
"""
query = f'''
DO $$
BEGIN
CREATE ROLE {database_user} WITH
PASSWORD '{database_password}'
NOSUPERUSER NOCREATEDB NOCREATEROLE INHERIT LOGIN NOREPLICATION
NOBYPASSRLS;
EXCEPTION WHEN duplicate_object THEN
ALTER ROLE {database_user} WITH
PASSWORD '{database_password}';
END
$$;'''
run_query(query)
def _drop_user(database_user: str):
"""Remove a user account with given username."""
run_query(f'DROP ROLE {database_user};')
def create_database(database_name: str, database_user: str,
database_password: str):
"""Create a new database and a user account to access it.
Database server is temporarily started if it is not running.
"""
query = f'''
CREATE EXTENSION IF NOT EXISTS dblink;
DO $$
BEGIN
PERFORM dblink_exec('',
'CREATE DATABASE {database_name} WITH OWNER {database_user}');
EXCEPTION WHEN duplicate_database THEN
ALTER DATABASE {database_name}
OWNER TO {database_user};
END
$$;'''
with action_utils.service_ensure_running('postgresql'):
_create_user(database_user, database_password)
run_query(query)
def drop_database(database_name: str, database_user: str):
"""Delete the database and the user account owning it.
Database server is temporarily started if it is not running.
"""
query = f'DROP DATABASE {database_name};'
with action_utils.service_ensure_running('postgresql'):
run_query(query)
_drop_user(database_user)
def dump_database(backup_file: str | pathlib.Path, database_name: str):
"""Dump PostgreSQL database to a file.
Database server is temporarily started if it is not running. Overwrite
file if it exists.
"""
backup_path = pathlib.Path(backup_file)
with action_utils.service_ensure_running('postgresql'):
with open(backup_path, 'w', encoding='utf-8') as file_handle:
_run_as([
'pg_dump', '--create', '--clean', '--if-exists', database_name
], stdout=file_handle)
def restore_database(backup_file: str | pathlib.Path, database_name: str,
database_user: str, database_password: str):
"""Restore database from a file.
Database server is temporarily started if it is not running. User account
is removed and recreated if it already exists. Drop database and recreate
if it already exists.
"""
with action_utils.service_ensure_running('postgresql'):
drop_database(database_name, database_user)
create_database(database_name, database_user, database_password)
with open(backup_file, 'r', encoding='utf-8') as file_handle:
_run_as(['psql', '--dbname', database_name], stdin=file_handle)
|
