Add functions for cloning postgres DB

rattail_fabric2/postgresql.py

@@ -26,6 +26,8 @@ Fabric Library for PostgreSQL
 
 from __future__ import unicode_literals, absolute_import
 
+import os
+
 from rattail_fabric2 import apt
 
 
@@ -94,3 +96,64 @@ def create_db(c, name, owner=None, port=None, checkfirst=True):
             owner='--owner={}'.format(owner) if owner else '',
             name=name)
         c.sudo(cmd)
+
+
+def drop_db(c, name, checkfirst=True):
+    """
+    Drop a PostgreSQL database.
+    """
+    if not checkfirst or db_exists(c, name):
+        c.sudo('sudo -u postgres dropdb {}'.format(name))
+
+
+def download_db(c, name, destination=None, port=None, exclude_tables=None):
+    """
+    Download a database from the server represented by ``c`` param.
+    """
+    if destination is None:
+        destination = './{}.sql.gz'.format(name)
+    c.run('touch {}.sql'.format(name))
+    c.run('chmod 0666 {}.sql'.format(name))
+    c.sudo('sudo -u postgres pg_dump {port} {exclude_tables} --file={name}.sql {name}'.format(
+        name=name,
+        port='--port={}'.format(port) if port else '',
+        exclude_tables='--exclude-table-data={}'.format(exclude_tables) if exclude_tables else '',
+    ))
+    c.run('gzip --force {}.sql'.format(name))
+    c.get('{}.sql.gz'.format(name), destination)
+    c.run('rm {}.sql.gz'.format(name))
+
+
+def clone_db(c, name, owner, download, user='rattail', force=False, workdir=None):
+    """
+    Clone a database from a (presumably live) server
+
+    :param name: Name of the database.
+
+    :param owner: Username of the user who is to own the database.
+
+    :param force: Whether the target database should be forcibly dropped, if it
+       exists already.
+    """
+    if db_exists(c, name):
+       if force:
+           drop_db(c, name, checkfirst=False)
+       else:
+           raise RuntimeError("Database '{}' already exists!".format(name))
+
+    create_db(c, name, owner=owner, checkfirst=False)
+
+    # upload database dump to target server
+    if workdir:
+        curdir = os.getcwd()
+        os.chdir(workdir)
+    download(c, '{}.sql.gz'.format(name), user=user)
+    c.put('{}.sql.gz'.format(name))
+    c.local('rm {}.sql.gz'.format(name))
+    if workdir:
+        os.chdir(curdir)
+
+    # restore database on target server
+    c.run('gunzip --force {}.sql.gz'.format(name))
+    c.sudo('sudo -u postgres psql --echo-errors --file={0}.sql {0}'.format(name))
+    c.run('rm {}.sql'.format(name))