From 306dbe65bf85169bcda087f850c847dd3f142b61 Mon Sep 17 00:00:00 2001
From: Ezra Pagel <ezra.pagel@icloud.com>
Date: Wed, 8 Nov 2023 09:55:03 -0600
Subject: [PATCH] add support for optional comments on statements

---
 README.md                  |  3 +++
 requirements.txt           |  2 +-
 scripts/flexible_freeze.py | 36 +++++++++++++++++++++++++++++-------
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index ee2f7fe..5d64a40 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,7 @@ Requires psycopg2
                               [--exclude-table-in-database EXCLUDE_TABLE_IN_DATABASE] [--no-freeze] [--no-analyze] [--vacuum]
                               [--pause PAUSE_TIME] [--freezeage FREEZEAGE] [--costdelay COSTDELAY] [--costlimit COSTLIMIT] [-t]
                               [--enforce-time] [-l LOGFILE] [-v] [--debug] [-U DBUSER] [-H DBHOST] [-p DBPORT] [-w DBPASS] [-st TABLE]
+                              [--comment COMMENT]
 
     optional arguments:
       -h, --help            show this help message and exit
@@ -61,6 +62,8 @@ Requires psycopg2
                             database password
       -st TABLE, --table TABLE
                             only process specified table
+      -c COMMENT, --comment COMMMENT
+                            specify an optional comment to annotate the query with
 
 Notes:
 
diff --git a/requirements.txt b/requirements.txt
index e049f93..4c9d9de 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1 @@
-psycopg2-binary==2.7.4
+psycopg2-binary>=2.8
diff --git a/scripts/flexible_freeze.py b/scripts/flexible_freeze.py
index 91725d7..c787b19 100755
--- a/scripts/flexible_freeze.py
+++ b/scripts/flexible_freeze.py
@@ -77,6 +77,8 @@ def timestamp():
                   help="database password")
 parser.add_argument("-st", "--table", dest="table",
                   help="only process specified  table", default=False)
+parser.add_argument("-c", "--comment", dest="comment",
+                  help="optional comment to annotate commands with", default=None)
 
 args = parser.parse_args()
 
@@ -121,6 +123,25 @@ def dbconnect(dbname, dbuser, dbhost, dbport, dbpass):
 
     return conn
 
+
+comment_start = "/*"
+comment_end = "*/"
+
+
+def build_statement(sql):
+    if args.comment is None:
+        return sql
+
+    comment = args.comment.replace(comment_start, "").replace(comment_end, "")
+
+    if "/*" in sql or "*/" in sql:
+        debug_print(f"sql statement already has comment, not appending: {sql}, {comment}")
+        return sql
+
+    sql += f" /* {comment} */"
+    debug_print(f"appended comment to sql: {sql}")
+    return sql
+
 def signal_handler(signal, frame):
     _print('exiting due to user interrupt')
     if conn:
@@ -191,9 +212,10 @@ def signal_handler(signal, frame):
         sys.exit(1)
 
     cur = conn.cursor()
-    cur.execute("""SELECT datname FROM pg_database
+    stmt = build_statement("""SELECT datname FROM pg_database
         WHERE datname NOT IN ('postgres','template1','template0')
         ORDER BY age(datfrozenxid) DESC""")
+    cur.execute(stmt)
     dblist = []
     for dbname in cur:
         dblist.append(dbname[0])
@@ -229,8 +251,8 @@ def signal_handler(signal, frame):
         continue
 
     cur = conn.cursor()
-    cur.execute("SET vacuum_cost_delay = {0}".format(args.costdelay))
-    cur.execute("SET vacuum_cost_limit = {0}".format(args.costlimit))
+    cur.execute(build_statement("SET vacuum_cost_delay = {0}".format(args.costdelay)))
+    cur.execute(build_statement("SET vacuum_cost_limit = {0}".format(args.costlimit)))
     
     # if vacuuming, get list of top tables to vacuum
     if args.skip_freeze:
@@ -271,7 +293,7 @@ def signal_handler(signal, frame):
             ORDER BY freeze_age DESC, table_bytes DESC
             LIMIT 1000;""".format(args.freezeage, args.minsizemb)
 
-    cur.execute(tabquery)
+    cur.execute(build_statement(tabquery))
     verbose_print("getting list of tables")
 
     table_resultset = cur.fetchall()
@@ -322,11 +344,11 @@ def signal_handler(signal, frame):
 
         try:
             if args.enforcetime:
-                excur.execute(timeout_query)
+                excur.execute(build_statement(timeout_query))
             else:
-                excur.execute("SET statement_timeout = 0")
+                excur.execute(build_statement("SET statement_timeout = 0"))
                 
-            excur.execute(exquery)
+            excur.execute(build_statement(exquery))
         except Exception as ex:
             _print("VACUUMing %s failed." % table)
             _print(str(ex))