Coverage for .tox/coverage/lib/python3.13/site-packages/wuttamess/postgres.py: 100%
34 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-20 15:16 -0500
1# -*- coding: utf-8; -*-
2################################################################################
3#
4# WuttaMess -- Fabric Automation Helpers
5# Copyright © 2024-2025 Lance Edgar
6#
7# This file is part of Wutta Framework.
8#
9# Wutta Framework is free software: you can redistribute it and/or modify it
10# under the terms of the GNU General Public License as published by the Free
11# Software Foundation, either version 3 of the License, or (at your option) any
12# later version.
13#
14# Wutta Framework is distributed in the hope that it will be useful, but
15# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
16# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
17# more details.
18#
19# You should have received a copy of the GNU General Public License along with
20# Wutta Framework. If not, see <http://www.gnu.org/licenses/>.
21#
22################################################################################
23"""
24PostgreSQL DB utilities
25"""
28def sql(c, sql_, database="", port=None, **kwargs):
29 """
30 Execute some SQL as the ``postgres`` user.
32 :param c: Fabric connection.
34 :param sql_: SQL string to execute.
36 :param database: Name of the database on which to execute the SQL.
37 If not specified, default ``postgres`` is assumed.
39 :param port: Optional port for PostgreSQL; default is 5432.
41 :returns: Result of the ``c.sudo()`` call (``psql -c`` command).
42 """
43 port = f" --port={port}" if port else ""
44 return c.sudo(
45 f'psql{port} --tuples-only --no-align --command="{sql_}" {database}',
46 user="postgres",
47 **kwargs,
48 )
51def user_exists(c, name, port=None):
52 """
53 Determine if a given PostgreSQL user exists.
55 :param c: Fabric connection.
57 :param name: Username to check for.
59 :param port: Optional port for PostgreSQL; default is 5432.
61 :returns: ``True`` if user exists, else ``False``.
62 """
63 user = sql(
64 c, f"SELECT rolname FROM pg_roles WHERE rolname = '{name}'", port=port
65 ).stdout.strip()
66 return bool(user)
69def create_user(c, name, password=None, port=None, checkfirst=True):
70 """
71 Create a PostgreSQL user account.
73 :param c: Fabric connection.
75 :param name: Username to create.
77 :param password: Optional password for the new user. If set, will
78 call :func:`set_user_password()`.
80 :param port: Optional port for PostgreSQL; default is 5432.
82 :param checkfirst: If true (the default), first call
83 :func:`user_exists()` and skip creating if already present. If
84 false, then try to create user with no check.
85 """
86 if not checkfirst or not user_exists(c, name, port=port):
87 portarg = f" --port={port}" if port else ""
88 c.sudo(
89 f"createuser{portarg} --no-createrole --no-superuser {name}",
90 user="postgres",
91 )
92 if password:
93 set_user_password(c, name, password, port=port)
96def set_user_password(c, name, password, port=None):
97 """
98 Set the password for a PostgreSQL user account.
100 :param c: Fabric connection.
102 :param name: Username whose password is to be set.
104 :param password: New password for the user.
106 :param port: Optional port for PostgreSQL; default is 5432.
107 """
108 sql(
109 c,
110 f"ALTER USER \\\"{name}\\\" PASSWORD '{password}';",
111 port=port,
112 hide=True,
113 echo=False,
114 )
117def db_exists(c, name, port=None):
118 """
119 Determine if a given PostgreSQL database exists.
121 :param c: Fabric connection.
123 :param name: Name of the database to check for.
125 :param port: Optional port for PostgreSQL; default is 5432.
127 :returns: ``True`` if database exists, else ``False``.
128 """
129 db = sql(
130 c, f"SELECT datname FROM pg_database WHERE datname = '{name}'", port=port
131 ).stdout.strip()
132 return db == name
135def create_db(c, name, owner=None, port=None, checkfirst=True):
136 """
137 Create a PostgreSQL database.
139 :param c: Fabric connection.
141 :param name: Name of the database to create.
143 :param owner: Optional role name to set as owner for the database.
145 :param port: Optional port for PostgreSQL; default is 5432.
147 :param checkfirst: If true (the default), first call
148 :func:`db_exists()` and skip creating if already present. If
149 false, then try to create DB with no check.
150 """
151 if not checkfirst or not db_exists(c, name, port=port):
152 port = f" --port={port}" if port else ""
153 owner = f" --owner={owner}" if owner else ""
154 c.sudo(f"createdb{port}{owner} {name}", user="postgres")
157def drop_db(c, name, checkfirst=True):
158 """
159 Drop a PostgreSQL database.
161 :param c: Fabric connection.
163 :param name: Name of the database to drop.
165 :param checkfirst: If true (the default), first call
166 :func:`db_exists()` and skip dropping if not present. If
167 false, then try to drop DB with no check.
168 """
169 if not checkfirst or db_exists(c, name):
170 c.sudo(f"dropdb {name}", user="postgres")
173def dump_db(c, name):
174 """
175 Dump a PostgreSQL database to file.
177 This uses the ``pg_dump`` and ``gzip`` commands to produce a
178 compressed SQL dump. The filename returned is based on the
179 ``name`` provided, e.g. ``mydbname.sql.gz``.
181 :param c: Fabric connection.
183 :param name: Name of the database to dump.
185 :returns: Base name of the output file. We only return the
186 filename and not the path, since the file is expected to exist
187 in the connected user's home folder.
188 """
189 sql_name = f"{name}.sql"
190 gz_name = f"{sql_name}.gz"
191 tmp_name = f"/tmp/{gz_name}"
193 # TODO: when pg_dump fails the command still succeeds! (would this work?)
194 # cmd = f'set -e && pg_dump {name} | gzip -c > {tmp_name}'
195 cmd = f"pg_dump {name} | gzip -c > {tmp_name}"
197 c.sudo(cmd, user="postgres")
198 c.run(f"cp {tmp_name} {gz_name}")
199 c.run(f"rm {tmp_name}")
201 return gz_name