# -*- coding: utf-8; -*-
################################################################################
#
# WuttaMess -- Fabric Automation Helpers
# Copyright © 2024-2025 Lance Edgar
#
# This file is part of Wutta Framework.
#
# Wutta Framework is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# Wutta Framework is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# Wutta Framework. If not, see <http://www.gnu.org/licenses/>.
#
################################################################################
"""
PostgreSQL DB utilities
"""
[docs]
def sql(c, sql_, database="", port=None, **kwargs):
"""
Execute some SQL as the ``postgres`` user.
:param c: Fabric connection.
:param sql_: SQL string to execute.
:param database: Name of the database on which to execute the SQL.
If not specified, default ``postgres`` is assumed.
:param port: Optional port for PostgreSQL; default is 5432.
"""
port = f" --port={port}" if port else ""
return c.sudo(
f'psql{port} --tuples-only --no-align --command="{sql_}" {database}',
user="postgres",
**kwargs,
)
[docs]
def user_exists(c, name, port=None):
"""
Determine if a given PostgreSQL user exists.
:param c: Fabric connection.
:param name: Username to check for.
:param port: Optional port for PostgreSQL; default is 5432.
:returns: ``True`` if user exists, else ``False``.
"""
user = sql(
c, f"SELECT rolname FROM pg_roles WHERE rolname = '{name}'", port=port
).stdout.strip()
return bool(user)
[docs]
def create_user(c, name, password=None, port=None, checkfirst=True):
"""
Create a PostgreSQL user account.
:param c: Fabric connection.
:param name: Username to create.
:param password: Optional password for the new user. If set, will
call :func:`set_user_password()`.
:param port: Optional port for PostgreSQL; default is 5432.
:param checkfirst: If true (the default), first check if user
exists and skip creating if already present. If false, then
try to create user with no check.
"""
if not checkfirst or not user_exists(c, name, port=port):
portarg = f" --port={port}" if port else ""
c.sudo(
f"createuser{portarg} --no-createrole --no-superuser {name}",
user="postgres",
)
if password:
set_user_password(c, name, password, port=port)
[docs]
def set_user_password(c, name, password, port=None):
"""
Set the password for a PostgreSQL user account.
:param c: Fabric connection.
:param name: Username whose password is to be set.
:param password: Password for the new user.
:param port: Optional port for PostgreSQL; default is 5432.
"""
sql(
c,
f"ALTER USER \\\"{name}\\\" PASSWORD '{password}';",
port=port,
hide=True,
echo=False,
)
[docs]
def db_exists(c, name, port=None):
"""
Determine if a given PostgreSQL database exists.
:param c: Fabric connection.
:param name: Name of the database to check for.
:param port: Optional port for PostgreSQL; default is 5432.
:returns: ``True`` if database exists, else ``False``.
"""
db = sql(
c, f"SELECT datname FROM pg_database WHERE datname = '{name}'", port=port
).stdout.strip()
return db == name
[docs]
def create_db(c, name, owner=None, port=None, checkfirst=True):
"""
Create a PostgreSQL database.
:param c: Fabric connection.
:param name: Name of the database to create.
:param owner: Optional role name to set as owner for the database.
:param port: Optional port for PostgreSQL; default is 5432.
:param checkfirst: If true (the default), first check if DB exists
and skip creating if already present. If false, then try to
create DB with no check.
"""
if not checkfirst or not db_exists(c, name, port=port):
port = f" --port={port}" if port else ""
owner = f" --owner={owner}" if owner else ""
c.sudo(f"createdb{port}{owner} {name}", user="postgres")
[docs]
def drop_db(c, name, checkfirst=True):
"""
Drop a PostgreSQL database.
:param c: Fabric connection.
:param name: Name of the database to drop.
:param checkfirst: If true (the default), first check if DB exists
and skip dropping if not present. If false, then try to drop
DB with no check.
"""
if not checkfirst or db_exists(c, name):
c.sudo(f"dropdb {name}", user="postgres")
[docs]
def dump_db(c, name):
"""
Dump a PostgreSQL database to file.
This uses the ``pg_dump`` and ``gzip`` commands to produce a
compressed SQL dump. The filename returned is based on the
``name`` provided, e.g. ``mydbname.sql.gz``.
:param c: Fabric connection.
:param name: Name of the database to dump.
:returns: Base name of the output file. We only return the
filename and not the path, since the file is expected to exist
in the connected user's home folder.
"""
sql_name = f"{name}.sql"
gz_name = f"{sql_name}.gz"
tmp_name = f"/tmp/{gz_name}"
# TODO: when pg_dump fails the command still succeeds! (would this work?)
# cmd = f'set -e && pg_dump {name} | gzip -c > {tmp_name}'
cmd = f"pg_dump {name} | gzip -c > {tmp_name}"
c.sudo(cmd, user="postgres")
c.run(f"cp {tmp_name} {gz_name}")
c.run(f"rm {tmp_name}")
return gz_name