feature: add support for optional comments on statements

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:
 

requirements.txt

@@ -1 +1 @@
-psycopg2-binary==2.7.4
+psycopg2-binary>=2.8

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))