Split up the dump/restore DB logic for postgresql

so that i can swap out portions of that in order to use rsync for big DBs

rattail_fabric2/postgresql.py

@@ -2,7 +2,7 @@
 ################################################################################
 #
 #  Rattail -- Retail Software Framework
-#  Copyright © 2010-2019 Lance Edgar
+#  Copyright © 2010-2020 Lance Edgar
 #
 #  This file is part of Rattail.
 #
@@ -137,12 +137,14 @@ def drop_db(c, name, checkfirst=True):
         c.sudo('dropdb {}'.format(name), user='postgres')
 
 
-def download_db(c, name, destination=None, port=None, exclude_tables=None):
+def dump_db(c, name, port=None, exclude_tables=None):
     """
-    Download a database from the server represented by ``c`` param.
+    Dump a database to file, on the server represented by ``c`` param.
+
+    This function returns the name of the DB dump file.  The name will not have
+    a path component as it's assumed to be in the home folder of the connection
+    user.
     """
-    if destination is None:
-        destination = './{}.sql.gz'.format(name)
     c.run('touch {}.sql'.format(name))
     c.run('chmod 0666 {}.sql'.format(name))
     cmd = 'pg_dump {port} {exclude_tables} --file={name}.sql {name}'.format(
@@ -151,8 +153,18 @@ def download_db(c, name, destination=None, port=None, exclude_tables=None):
         exclude_tables='--exclude-table-data={}'.format(exclude_tables) if exclude_tables else '')
     c.sudo(cmd, user='postgres')
     c.run('gzip --force {}.sql'.format(name))
-    c.get('{}.sql.gz'.format(name), destination)
+    return '{}.sql.gz'.format(name)
-    c.run('rm {}.sql.gz'.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)
+    dumpfile = dump_db(c, name, port=port, exclude_tables=exclude_tables)
+    c.get(dumpfile, destination)
+    c.run('rm {}'.format(dumpfile))
 
 
 def clone_db(c, name, owner, download, user='rattail', force=False, workdir=None):
@@ -185,6 +197,22 @@ def clone_db(c, name, owner, download, user='rattail', force=False, workdir=None
         os.chdir(curdir)
 
     # restore database on target server
-    c.run('gunzip --force {}.sql.gz'.format(name))
+    restore_db(c, name, '{}.sql.gz'.format(name))
-    c.sudo('psql --echo-errors --file={0}.sql {0}'.format(name), user='postgres')
+
-    c.run('rm {}.sql'.format(name))
+
+def restore_db(c, name, path):
+    """
+    Restore a database from a dump file.
+
+    :param name: Name of the database to restore.
+
+    :param path: Path to the DB dump file, which should end in ``.sql.gz``
+    """
+    if not path.endswith('.sql.gz'):
+        raise ValueError("Path to dump file must end in `.sql.gz`")
+
+    c.sudo('gunzip --force {}'.format(path))
+    sql_path = path[:-3]
+    c.sudo('psql --echo-errors --file={} {}'.format(sql_path, name),
+           user='postgres')
+    c.sudo('rm {}'.format(sql_path))