source: rattail-fabric/rattail_fabric/postgresql.py @ c87ef42

Last change on this file since c87ef42 was c87ef42, checked in by Lance Edgar <ledgar@…>, 18 months ago

Fix how we run sudo commands as postgres user

per similar changes made in rattail-fabric2

  • Property mode set to 100644
File size: 6.3 KB
Line 
1# -*- coding: utf-8; -*-
2################################################################################
3#
4#  Rattail -- Retail Software Framework
5#  Copyright © 2010-2018 Lance Edgar
6#
7#  This file is part of Rattail.
8#
9#  Rattail is free software: you can redistribute it and/or modify it under the
10#  terms of the GNU General Public License as published by the Free Software
11#  Foundation, either version 3 of the License, or (at your option) any later
12#  version.
13#
14#  Rattail is distributed in the hope that it will be useful, but WITHOUT ANY
15#  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
16#  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
17#  details.
18#
19#  You should have received a copy of the GNU General Public License along with
20#  Rattail.  If not, see <http://www.gnu.org/licenses/>.
21#
22################################################################################
23"""
24Fabric Library for PostgreSQL
25"""
26
27from __future__ import unicode_literals, absolute_import
28
29import os
30import re
31
32from fabric.api import sudo, run, get, hide, abort, put, local
33
34from rattail_fabric import apt
35
36
37def install():
38    """
39    Install the PostgreSQL database service
40    """
41    apt.install('postgresql')
42
43
44def get_version():
45    """
46    Fetch the version of PostgreSQL running on the target system
47    """
48    result = sudo('psql --version')
49    if result.succeeded:
50        match = re.match(r'^psql \(PostgreSQL\) (\d+\.\d+)(?:\.\d+)?', result)
51        if match:
52            return float(match.group(1))
53
54
55def sql(sql, database='', port=None):
56    """
57    Execute some SQL as the 'postgres' user.
58    """
59    cmd = 'psql {port} --tuples-only --no-align --command="{sql}" {database}'.format(
60        port='--port={}'.format(port) if port else '',
61        sql=sql, database=database)
62    return sudo(cmd, user='postgres')
63
64
65def script(path, database='', port=None, user=None, password=None):
66    """
67    Execute a SQL script.  By default this will run as 'postgres' user, but can
68    use PGPASSWORD authentication if necessary.
69    """
70    port = '--port={}'.format(port) if port else ''
71    if user and password:
72        with hide('running'):
73            kw = dict(pw=password, user=user, port=port, path=path, db=database)
74            return sudo(" PGPASSWORD='{pw}' psql --host=localhost {port} --username='{user}' --file='{path}' {db}".format(**kw))
75
76    else: # run as postgres
77        kw = dict(port=port, path=path, db=database)
78        return sudo("psql {port} --file='{path}' {db}".format(**kw), user='postgres')
79
80
81def user_exists(name, port=None):
82    """
83    Determine if a given PostgreSQL user exists.
84    """
85    user = sql("SELECT rolname FROM pg_roles WHERE rolname = '{0}'".format(name), port=port)
86    return bool(user)
87
88
89def create_user(name, password=None, port=None, checkfirst=True, createdb=False):
90    """
91    Create a PostgreSQL user account.
92    """
93    if not checkfirst or not user_exists(name, port=port):
94        cmd = 'createuser {port} {createdb} --no-createrole --no-superuser {name}'.format(
95            port='--port={}'.format(port) if port else '',
96            createdb='--{}createdb'.format('' if createdb else 'no-'),
97            name=name)
98        sudo(cmd, user='postgres')
99        if password:
100            set_user_password(name, password, port=port)
101
102
103def set_user_password(name, password, port=None):
104    """
105    Set the password for a PostgreSQL user account
106    """
107    with hide('running'):
108        sql("ALTER USER \\\"{}\\\" PASSWORD '{}';".format(name, password), port=port)
109
110
111def db_exists(name, port=None):
112    """
113    Determine if a given PostgreSQL database exists.
114    """
115    db = sql("SELECT datname FROM pg_database WHERE datname = '{0}'".format(name), port=port)
116    return db == name
117
118
119def create_db(name, owner=None, port=None, checkfirst=True):
120    """
121    Create a PostgreSQL database.
122    """
123    if not checkfirst or not db_exists(name, port=port):
124        cmd = 'createdb {port} {owner} {name}'.format(
125            port='--port={}'.format(port) if port else '',
126            owner='--owner={}'.format(owner) if owner else '',
127            name=name)
128        sudo(cmd, user='postgres')
129
130
131def create_schema(name, dbname, owner='rattail', port=None):
132    """
133    Create a schema within a PostgreSQL database.
134    """
135    sql_ = "create schema if not exists {} authorization {}".format(name, owner)
136    sql(sql_, database=dbname, port=port)
137
138
139def drop_db(name, checkfirst=True):
140    """
141    Drop a PostgreSQL database.
142    """
143    if not checkfirst or db_exists(name):
144        sudo('dropdb {}'.format(name), user='postgres')
145
146
147def download_db(name, destination=None, port=None, exclude_tables=None):
148    """
149    Download a database from the "current" server.
150    """
151    if destination is None:
152        destination = './{0}.sql.gz'.format(name)
153    run('touch {0}.sql'.format(name))
154    run('chmod 0666 {0}.sql'.format(name))
155    cmd = 'pg_dump {port} {exclude_tables} --file={name}.sql {name}'.format(
156        name=name,
157        port='--port={}'.format(port) if port else '',
158        exclude_tables='--exclude-table-data={}'.format(exclude_tables) if exclude_tables else '')
159    sudo(cmd, user='postgres')
160    run('gzip --force {0}.sql'.format(name))
161    get('{0}.sql.gz'.format(name), destination)
162    run('rm {0}.sql.gz'.format(name))
163
164
165def clone_db(name, owner, download, user='rattail', force=False, workdir=None):
166    """
167    Clone a database from a (presumably live) server
168
169    :param name: Name of the database.
170
171    :param owner: Username of the user who is to own the database.
172
173    :param force: Whether the target database should be forcibly dropped, if it
174       exists already.
175    """
176    if db_exists(name):
177       if force:
178           drop_db(name, checkfirst=False)
179       else:
180           abort("Database '{}' already exists!".format(name))
181
182    create_db(name, owner=owner, checkfirst=False)
183
184    # upload database dump to target server
185    if workdir:
186        curdir = os.getcwd()
187        os.chdir(workdir)
188    download('{}.sql.gz'.format(name), user=user)
189    put('{}.sql.gz'.format(name))
190    local('rm {}.sql.gz'.format(name))
191    if workdir:
192        os.chdir(curdir)
193
194    # restore database on target server
195    run('gunzip --force {}.sql.gz'.format(name))
196    sudo('psql --echo-errors --file={0}.sql {0}'.format(name), user='postgres')
197    run('rm {}.sql'.format(name))
Note: See TracBrowser for help on using the repository browser.