Merge ~cjwatson/launchpad:update-pgkill-scripts into launchpad:master

[Colin Watson]

Update pgkillactive and pgkillidle scripts

Similar versions of these scripts exist in lp:losa-db-scripts and lp:postgresql-charm, and have been a bit more actively maintained there. However, it's still useful to have copies of these in the Launchpad tree, since the simplest way to migrate some current cron jobs from production will be to run these as part of the `launchpad-admin` charm.

Sync our versions up more closely with those maintained elsewhere, making the following changes:

 * Exclude the current process from being killed.
 * Terminate the PostgreSQL backend process remotely using `pg_terminate_backend`, rather than relying on running on the database host itself.
 * Make the `pgkillidle -u` option match either the username or the application name.
 * Update `pgkillidle` for `pg_stat_activity` changes in PostgreSQL >= 9.2.
 * Make `pgkillidle` print a little more information on the process being killed.
 * Minor changes to docstrings and help text.

[Ines Almeida (ines-almeida)]

Looks good

review: Approve

diff --git a/utilities/pgkillactive.py b/utilities/pgkillactive.py
index c8e5e44..803e3d1 100755
--- a/utilities/pgkillactive.py
+++ b/utilities/pgkillactive.py
@@ -3,15 +3,12 @@
 # Copyright 2009 Canonical Ltd.  This software is licensed under the
 # GNU Affero General Public License version 3 (see the file LICENSE).
 
-"""Kill transaction that have hung around for too long.
-"""
+"""Kill transactions that have hung around for too long."""
 
 __all__ = []
 
 import _pythonpath  # noqa: F401
 
-import os
-import signal
 import sys
 from optparse import OptionParser
 
@@ -34,7 +31,10 @@ def main():
         type="int",
         dest="max_seconds",
         default=60 * 60,
-        help="Maximum seconds time connections are allowed to remain active.",
+        help=(
+            "Connections with a transaction older than MAX_SECONDS seconds "
+            "will be killed. If 0, all matched connections will be killed."
+        ),
     )
     parser.add_option(
         "-q",
@@ -76,7 +76,10 @@ def main():
         """
         SELECT usename, pid, backend_start, xact_start
         FROM pg_stat_activity
-        WHERE xact_start < CURRENT_TIMESTAMP - '%d seconds'::interval %s
+        WHERE
+            pid <> pg_backend_pid()
+            AND xact_start < CURRENT_TIMESTAMP - '%d seconds'::interval
+            %s
         ORDER BY pid
         """
         % (options.max_seconds, user_match_sql),
@@ -88,7 +91,7 @@ def main():
     if len(rows) == 0:
         if not options.quiet:
             print("No transactions to kill")
-            return 0
+        return 0
 
     for usename, pid, backend_start, transaction_start in rows:
         print(
@@ -101,7 +104,8 @@ def main():
             )
         )
         if not options.dry_run:
-            os.kill(pid, signal.SIGTERM)
+            cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
+    cur.close()
     return 0
 
 
diff --git a/utilities/pgkillidle.py b/utilities/pgkillidle.py
index 9985ae8..6b6df3b 100755
--- a/utilities/pgkillidle.py
+++ b/utilities/pgkillidle.py
@@ -3,15 +3,12 @@
 # Copyright 2009 Canonical Ltd.  This software is licensed under the
 # GNU Affero General Public License version 3 (see the file LICENSE).
 
-"""Kill <IDLE> in transaction connections that have hung around for too long.
-"""
+"""Kill idle-in-transaction connections that have hung around for too long."""
 
 __all__ = []
 
 import _pythonpath  # noqa: F401
 
-import os
-import signal
 import sys
 from optparse import OptionParser
 
@@ -64,17 +61,24 @@ def main():
     if len(args) > 0:
         parser.error("Too many arguments")
 
-    ignore_sql = " AND usename <> %s" * len(options.ignore or [])
+    ignore_sql = " AND %s NOT IN (usename, application_name)" * len(
+        options.ignore or []
+    )
 
     con = psycopg2.connect(options.connect_string)
     cur = con.cursor()
     cur.execute(
         """
-        SELECT usename, pid, backend_start, query_start
+        SELECT
+            usename, application_name, datname, pid,
+            backend_start, state_change, AGE(NOW(), state_change) AS age
         FROM pg_stat_activity
-        WHERE query = '<IDLE> in transaction'
-            AND query_start < CURRENT_TIMESTAMP - '%d seconds'::interval %s
-        ORDER BY pid
+        WHERE
+            pid <> pg_backend_pid()
+            AND state = 'idle in transaction'
+            AND state_change < CURRENT_TIMESTAMP - '%d seconds'::interval
+            %s
+        ORDER BY age
         """
         % (options.max_idle_seconds, ignore_sql),
         options.ignore,
@@ -85,20 +89,17 @@ def main():
     if len(rows) == 0:
         if not options.quiet:
             print("No IDLE transactions to kill")
-            return 0
+        return 0
 
-    for usename, pid, backend_start, query_start in rows:
-        print(
-            "Killing %s(%d), %s, %s"
-            % (
-                usename,
-                pid,
-                backend_start,
-                query_start,
-            )
-        )
+    for usename, appname, datname, pid, backend, state, age in rows:
+        print(80 * "=")
+        print("Killing %s(%d) %s from %s:" % (usename, pid, appname, datname))
+        print("    backend start: %s" % (backend,))
+        print("    idle start:    %s" % (state,))
+        print("    age:           %s" % (age,))
         if not options.dryrun:
-            os.kill(pid, signal.SIGTERM)
+            cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
+    cur.close()
     return 0