Coverage for .tox/coverage/lib/python3.11/site-packages/wuttamess/postgres.py: 100%
34 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-08 07:01 -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.
40 """
41 port = f" --port={port}" if port else ""
42 return c.sudo(
43 f'psql{port} --tuples-only --no-align --command="{sql_}" {database}',
44 user="postgres",
45 **kwargs,
46 )
49def user_exists(c, name, port=None):
50 """
51 Determine if a given PostgreSQL user exists.
53 :param c: Fabric connection.
55 :param name: Username to check for.
57 :param port: Optional port for PostgreSQL; default is 5432.
59 :returns: ``True`` if user exists, else ``False``.
60 """
61 user = sql(
62 c, f"SELECT rolname FROM pg_roles WHERE rolname = '{name}'", port=port
63 ).stdout.strip()
64 return bool(user)
67def create_user(c, name, password=None, port=None, checkfirst=True):
68 """
69 Create a PostgreSQL user account.
71 :param c: Fabric connection.
73 :param name: Username to create.
75 :param password: Optional password for the new user. If set, will
76 call :func:`set_user_password()`.
78 :param port: Optional port for PostgreSQL; default is 5432.
80 :param checkfirst: If true (the default), first check if user
81 exists and skip creating if already present. If false, then
82 try to create user with no check.
83 """
84 if not checkfirst or not user_exists(c, name, port=port):
85 portarg = f" --port={port}" if port else ""
86 c.sudo(
87 f"createuser{portarg} --no-createrole --no-superuser {name}",
88 user="postgres",
89 )
90 if password:
91 set_user_password(c, name, password, port=port)
94def set_user_password(c, name, password, port=None):
95 """
96 Set the password for a PostgreSQL user account.
98 :param c: Fabric connection.
100 :param name: Username whose password is to be set.
102 :param password: Password for the new user.
104 :param port: Optional port for PostgreSQL; default is 5432.
105 """
106 sql(
107 c,
108 f"ALTER USER \\\"{name}\\\" PASSWORD '{password}';",
109 port=port,
110 hide=True,
111 echo=False,
112 )
115def db_exists(c, name, port=None):
116 """
117 Determine if a given PostgreSQL database exists.
119 :param c: Fabric connection.
121 :param name: Name of the database to check for.
123 :param port: Optional port for PostgreSQL; default is 5432.
125 :returns: ``True`` if database exists, else ``False``.
126 """
127 db = sql(
128 c, f"SELECT datname FROM pg_database WHERE datname = '{name}'", port=port
129 ).stdout.strip()
130 return db == name
133def create_db(c, name, owner=None, port=None, checkfirst=True):
134 """
135 Create a PostgreSQL database.
137 :param c: Fabric connection.
139 :param name: Name of the database to create.
141 :param owner: Optional role name to set as owner for the database.
143 :param port: Optional port for PostgreSQL; default is 5432.
145 :param checkfirst: If true (the default), first check if DB exists
146 and skip creating if already present. If false, then try to
147 create DB with no check.
148 """
149 if not checkfirst or not db_exists(c, name, port=port):
150 port = f" --port={port}" if port else ""
151 owner = f" --owner={owner}" if owner else ""
152 c.sudo(f"createdb{port}{owner} {name}", user="postgres")
155def drop_db(c, name, checkfirst=True):
156 """
157 Drop a PostgreSQL database.
159 :param c: Fabric connection.
161 :param name: Name of the database to drop.
163 :param checkfirst: If true (the default), first check if DB exists
164 and skip dropping if not present. If false, then try to drop
165 DB with no check.
166 """
167 if not checkfirst or db_exists(c, name):
168 c.sudo(f"dropdb {name}", user="postgres")
171def dump_db(c, name):
172 """
173 Dump a PostgreSQL database to file.
175 This uses the ``pg_dump`` and ``gzip`` commands to produce a
176 compressed SQL dump. The filename returned is based on the
177 ``name`` provided, e.g. ``mydbname.sql.gz``.
179 :param c: Fabric connection.
181 :param name: Name of the database to dump.
183 :returns: Base name of the output file. We only return the
184 filename and not the path, since the file is expected to exist
185 in the connected user's home folder.
186 """
187 sql_name = f"{name}.sql"
188 gz_name = f"{sql_name}.gz"
189 tmp_name = f"/tmp/{gz_name}"
191 # TODO: when pg_dump fails the command still succeeds! (would this work?)
192 # cmd = f'set -e && pg_dump {name} | gzip -c > {tmp_name}'
193 cmd = f"pg_dump {name} | gzip -c > {tmp_name}"
195 c.sudo(cmd, user="postgres")
196 c.run(f"cp {tmp_name} {gz_name}")
197 c.run(f"rm {tmp_name}")
199 return gz_name