db: Add more utilities for managing PostgreSQL databases

- Add methods for creating/dropping user/database.

Tests:

- ttrss and miniflux functional tests work which check for backup/restore.

Signed-off-by: Sunil Mohan Adapa <sunil@medhas.org>
Reviewed-by: James Valleroy <jvalleroy@mailbox.org>

plinth/db/__init__.py

@@ -1,10 +1,6 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
-"""
-Common utilities to help with handling a database.
-"""
+"""Common utilities to help with handling a database."""
 
-import pathlib
-import subprocess
 import threading
 from typing import ClassVar
 
@@ -85,51 +81,3 @@ class DBLock:
 # most of the significant cases where we have seen database lock issues.
 
 lock = DBLock()
-
-
-#
-# PostgreSQL utilites
-#
-def _run_as_postgres(command, stdin=None, stdout=None):
-    """Run a command as postgres user."""
-    command = ['sudo', '--user', 'postgres'] + command
-    return subprocess.run(command, stdin=stdin, stdout=stdout, check=True)
-
-
-def postgres_dump_database(backup_file: str, database_name: str,
-                           database_user: str):
-    """Dump PostgreSQL database to a file.
-
-    Overwrites file if it exists. Uses pg_dump utility from postgres package
-    (needs to be installed).
-    """
-    backup_path = pathlib.Path(backup_file)
-    backup_path.parent.mkdir(parents=True, exist_ok=True)
-    with open(backup_path, 'w', encoding='utf-8') as file_handle:
-        process = _run_as_postgres(['pg_dumpall', '--roles-only'],
-                                   stdout=subprocess.PIPE)
-        file_handle.write(f'DROP ROLE IF EXISTS {database_user};\n')
-        for line in process.stdout.decode().splitlines():
-            if database_user in line:
-                file_handle.write(line + '\n')
-
-    with open(backup_path, 'a', encoding='utf-8') as file_handle:
-        _run_as_postgres(
-            ['pg_dump', '--create', '--clean', '--if-exists', database_name],
-            stdout=file_handle)
-
-
-def postgres_restore_database(backup_file: str, database_name):
-    """Restore PostgreSQL database from a file.
-
-    Drops database and recreates it. Uses pg_dump utility from postgres package
-    (needs to be installed).
-    """
-    # This is needed for old backups only. New backups include 'DROP DATABASE
-    # IF EXISTS' and 'CREATE DATABASE' statements.
-    _run_as_postgres(['dropdb', database_name])
-    _run_as_postgres(['createdb', database_name])
-
-    with open(backup_file, 'r', encoding='utf-8') as file_handle:
-        _run_as_postgres(['psql', '--dbname', database_name],
-                         stdin=file_handle)

plinth/db/postgres.py

@@ -0,0 +1,117 @@
+# 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
+import subprocess
+
+from plinth import action_utils
+
+
+def _run_as(command, **kwargs):
+    """Run a command as 'postgres' user."""
+    command = ['sudo', '--user', 'postgres'] + command
+    return subprocess.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)
+    backup_path.parent.mkdir(parents=True, exist_ok=True)
+    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)

plinth/modules/miniflux/privileged.py

@@ -10,8 +10,9 @@ from urllib.parse import urlparse
 
 import pexpect
 
-from plinth import action_utils, db
+from plinth import action_utils
 from plinth.actions import privileged, secret_str
+from plinth.db import postgres
 from plinth.utils import is_non_empty_file
 
 STATIC_SETTINGS = {
@@ -142,12 +143,12 @@ def _get_database_config():
 def dump_database():
     """Dump database to file."""
     config = _get_database_config()
-    db.postgres_dump_database(DB_BACKUP_FILE, config['database'],
-                              config['user'])
+    postgres.dump_database(DB_BACKUP_FILE, config['database'])
 
 
 @privileged
 def restore_database():
     """Restore database from file."""
     config = _get_database_config()
-    db.postgres_restore_database(DB_BACKUP_FILE, config['database'])
+    postgres.restore_database(DB_BACKUP_FILE, config['database'],
+                              config['user'], config['password'])

plinth/modules/ttrss/privileged.py

@@ -3,8 +3,9 @@
 
 import augeas
 
-from plinth import action_utils, db
+from plinth import action_utils
 from plinth.actions import privileged
+from plinth.db import postgres
 
 CONFIG_FILE = '/etc/tt-rss/config.php'
 DEFAULT_FILE = '/etc/default/tt-rss'
@@ -121,15 +122,15 @@ def enable_api_access():
 def dump_database():
     """Dump database to file."""
     config = _get_database_config()
-    db.postgres_dump_database(DB_BACKUP_FILE, config['database'],
-                              config['user'])
+    postgres.dump_database(DB_BACKUP_FILE, config['database'])
 
 
 @privileged
 def restore_database():
     """Restore database from file."""
     config = _get_database_config()
-    db.postgres_restore_database(DB_BACKUP_FILE, config['database'])
+    postgres.restore_database(DB_BACKUP_FILE, config['database'],
+                              config['user'], config['password'])
 
 
 def load_augeas():