Coverage for .tox/coverage/lib/python3.11/site-packages/wuttamess/postgres.py: 100%

1# -*- coding: utf-8; -*- 

2################################################################################ 

3# 

4# WuttaMess -- Fabric Automation Helpers 

5# Copyright © 2024 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""" 

26 

27 

28def sql(c, sql, database='', port=None, **kwargs): 

29 """ 

30 Execute some SQL as the ``postgres`` user. 

31 

32 :param c: Fabric connection. 

33 

34 :param sql: SQL string to execute. 

35 

36 :param database: Name of the database on which to execute the SQL. 

37 If not specified, default ``postgres`` is assumed. 

38 

39 :param port: Optional port for PostgreSQL; default is 5432. 

40 """ 

41 port = f' --port={port}' if port else '' 

42 return c.sudo(f'psql{port} --tuples-only --no-align --command="{sql}" {database}', 

43 user='postgres', **kwargs) 

44 

45 

46def user_exists(c, name, port=None): 

47 """ 

48 Determine if a given PostgreSQL user exists. 

49 

50 :param c: Fabric connection. 

51 

52 :param name: Username to check for. 

53 

54 :param port: Optional port for PostgreSQL; default is 5432. 

55 

56 :returns: ``True`` if user exists, else ``False``. 

57 """ 

58 user = sql(c, f"SELECT rolname FROM pg_roles WHERE rolname = '{name}'", port=port).stdout.strip() 

59 return bool(user) 

60 

61 

62def create_user(c, name, password=None, port=None, checkfirst=True): 

63 """ 

64 Create a PostgreSQL user account. 

65 

66 :param c: Fabric connection. 

67 

68 :param name: Username to create. 

69 

70 :param password: Optional password for the new user. If set, will 

71 call :func:`set_user_password()`. 

72 

73 :param port: Optional port for PostgreSQL; default is 5432. 

74 

75 :param checkfirst: If true (the default), first check if user 

76 exists and skip creating if already present. If false, then 

77 try to create user with no check. 

78 """ 

79 if not checkfirst or not user_exists(c, name, port=port): 

80 portarg = f' --port={port}' if port else '' 

81 c.sudo(f'createuser{portarg} --no-createrole --no-superuser {name}', 

82 user='postgres') 

83 if password: 

84 set_user_password(c, name, password, port=port) 

85 

86 

87def set_user_password(c, name, password, port=None): 

88 """ 

89 Set the password for a PostgreSQL user account. 

90 

91 :param c: Fabric connection. 

92 

93 :param name: Username whose password is to be set. 

94 

95 :param password: Password for the new user. 

96 

97 :param port: Optional port for PostgreSQL; default is 5432. 

98 """ 

99 sql(c, f"ALTER USER \\\"{name}\\\" PASSWORD '{password}';", port=port, hide=True, echo=False) 

100 

101 

102def db_exists(c, name, port=None): 

103 """ 

104 Determine if a given PostgreSQL database exists. 

105 

106 :param c: Fabric connection. 

107 

108 :param name: Name of the database to check for. 

109 

110 :param port: Optional port for PostgreSQL; default is 5432. 

111 

112 :returns: ``True`` if database exists, else ``False``. 

113 """ 

114 db = sql(c, f"SELECT datname FROM pg_database WHERE datname = '{name}'", port=port).stdout.strip() 

115 return db == name 

116 

117 

118def create_db(c, name, owner=None, port=None, checkfirst=True): 

119 """ 

120 Create a PostgreSQL database. 

121 

122 :param c: Fabric connection. 

123 

124 :param name: Name of the database to create. 

125 

126 :param owner: Optional role name to set as owner for the database. 

127 

128 :param port: Optional port for PostgreSQL; default is 5432. 

129 

130 :param checkfirst: If true (the default), first check if DB exists 

131 and skip creating if already present. If false, then try to 

132 create DB with no check. 

133 """ 

134 if not checkfirst or not db_exists(c, name, port=port): 

135 port = f' --port={port}' if port else '' 

136 owner = f' --owner={owner}' if owner else '' 

137 c.sudo(f'createdb{port}{owner} {name}', 

138 user='postgres') 

139 

140 

141def drop_db(c, name, checkfirst=True): 

142 """ 

143 Drop a PostgreSQL database. 

144 

145 :param c: Fabric connection. 

146 

147 :param name: Name of the database to drop. 

148 

149 :param checkfirst: If true (the default), first check if DB exists 

150 and skip dropping if not present. If false, then try to drop 

151 DB with no check. 

152 """ 

153 if not checkfirst or db_exists(c, name): 

154 c.sudo(f'dropdb {name}', user='postgres') 

155 

156 

157def dump_db(c, name): 

158 """ 

159 Dump a PostgreSQL database to file. 

160 

161 This uses the ``pg_dump`` and ``gzip`` commands to produce a 

162 compressed SQL dump. The filename returned is based on the 

163 ``name`` provided, e.g. ``mydbname.sql.gz``. 

164 

165 :param c: Fabric connection. 

166 

167 :param name: Name of the database to dump. 

168 

169 :returns: Base name of the output file. We only return the 

170 filename and not the path, since the file is expected to exist 

171 in the connected user's home folder. 

172 """ 

173 sql_name = f'{name}.sql' 

174 gz_name = f'{sql_name}.gz' 

175 tmp_name = f'/tmp/{gz_name}' 

176 

177 # TODO: when pg_dump fails the command still succeeds! (would this work?) 

178 #cmd = f'set -e && pg_dump {name} | gzip -c > {tmp_name}' 

179 cmd = f'pg_dump {name} | gzip -c > {tmp_name}' 

180 

181 c.sudo(cmd, user='postgres') 

182 c.run(f"cp {tmp_name} {gz_name}") 

183 c.run(f"rm {tmp_name}") 

184 

185 return gz_name