From 24d632b7e35b6015ddfd8f0d2ea409155650d06d Mon Sep 17 00:00:00 2001
From: Lance Edgar <lance@edbob.org>
Date: Sat, 16 Sep 2023 16:56:29 -0500
Subject: [PATCH] Add option to skip raw SQL file when dumping postgres DB

trying to cut down on disk space, we'll see how well this works..
---
 rattail_fabric2/postgresql.py | 40 ++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/rattail_fabric2/postgresql.py b/rattail_fabric2/postgresql.py
index f2d8440..19c4591 100644
--- a/rattail_fabric2/postgresql.py
+++ b/rattail_fabric2/postgresql.py
@@ -28,6 +28,7 @@ import os
 import re
 
 from rattail_fabric2 import apt, append, contains, sed, uncomment
+from rattail.util import shlex_join
 
 
 def install(c):
@@ -199,7 +200,8 @@ def drop_db(c, name, checkfirst=True):
         c.sudo('dropdb {}'.format(name), user='postgres')
 
 
-def dump_db(c, name, port=None, exclude_tables=None):
+def dump_db(c, name, port=None, exclude_tables=None,
+            skip_raw_file=False):
     """
     Dump a database to file, on the server represented by ``c`` param.
 
@@ -209,22 +211,40 @@ def dump_db(c, name, port=None, exclude_tables=None):
     """
     c.run('touch {}.sql'.format(name))
     c.run('chmod 0666 {}.sql'.format(name))
-    cmd = '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.sudo(cmd, user='postgres')
-    c.run('gzip --force {}.sql'.format(name))
-    return '{}.sql.gz'.format(name)
+
+    sql_name = f'{name}.sql'
+    gz_name = f'{sql_name}.gz'
+    filename = gz_name if skip_raw_file else sql_name
+
+    port = f'--port={port}' if port else ''
+    exclude_tables = f'--exclude-table-data={exclude_tables}' if exclude_tables else ''
+    filename = '' if skip_raw_file else f'--file={filename}'
+    cmd = f'pg_dump {port} {exclude_tables} {filename} {name}'
+
+    if skip_raw_file:
+        tmp_name = f'/tmp/{gz_name}'
+        cmd = f'{cmd} | gzip -c > {tmp_name}'
+        cmd = shlex_join(['bash', '-c', cmd])
+        c.sudo(cmd, user='postgres')
+        # TODO: should remove this file
+        c.run(f"cp {tmp_name} {gz_name}")
+
+    else:
+        c.sudo(cmd, user='postgres')
+        c.run(f'gzip --force {sql_name}')
+
+    return gz_name
 
 
-def download_db(c, name, destination=None, port=None, exclude_tables=None):
+def download_db(c, name, destination=None, port=None, exclude_tables=None,
+                skip_raw_file=False):
     """
     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)
+    dumpfile = dump_db(c, name, port=port, exclude_tables=exclude_tables,
+                       skip_raw_file=skip_raw_file)
     c.get(dumpfile, destination)
     c.run('rm {}'.format(dumpfile))