From dcb07c225ec461f0c852f8d6935fe93f742a3283 Mon Sep 17 00:00:00 2001
From: Julien Riou <julien.riou@outlook.com>
Date: Sat, 30 Jun 2018 11:11:24 +0200
Subject: [PATCH] Add cancel option to terminate queries

---
 README.md                |  1 +
 base/config.go           |  1 +
 base/db.go               | 14 ++++++++++++++
 cmd/pgterminate/main.go  |  1 +
 terminator/terminator.go |  6 +++++-
 5 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 8fa19dc..7e2f5ac 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,7 @@ With `pgterminate`, you shouldn't be paged at night because some queries has loc
 # Highlights
 * `pgterminate` name is derived from `pg_terminate_backend` function, it terminates backends.
 * backends are called sessions in `pgterminate`.
+* use `cancel` option to terminate query instead of session.
 * `active` sessions are backends in `active` state for more than `active-timeout` seconds.
 * `idle` sessions are backends in `idle`, `idle in transaction` or `idle in transaction (abort)` state for more than `idle-timeout` seconds.
 * at least one of `active-timeout` and `idle-timeout` parameter is required, both can be used.
diff --git a/base/config.go b/base/config.go
index 71f2921..ee2eeb1 100644
--- a/base/config.go
+++ b/base/config.go
@@ -38,6 +38,7 @@ type Config struct {
 	ExcludeUsers              StringFlags `yaml:"exclude-users"`
 	ExcludeUsersRegex         string      `yaml:"exclude-users-regex"`
 	ExcludeUsersRegexCompiled *regexp.Regexp
+	Cancel                    bool `yaml:"cancel"`
 }
 
 func init() {
diff --git a/base/db.go b/base/db.go
index afc3e29..e2b8737 100644
--- a/base/db.go
+++ b/base/db.go
@@ -84,3 +84,17 @@ func (db *Db) TerminateSessions(sessions []Session) {
 		Panic(err)
 	}
 }
+
+// CancelSessions terminates current query of a list of sessions
+func (db *Db) CancelSessions(sessions []Session) {
+	var pids []int64
+	for _, session := range sessions {
+		pids = append(pids, session.Pid)
+	}
+	if len(pids) > 0 {
+		query := `select pg_cancel_backend(pid) from pg_stat_activity where pid = any($1);`
+		log.Debugf("query: %s\n", query)
+		_, err := db.conn.Exec(query, pq.Array(pids))
+		Panic(err)
+	}
+}
diff --git a/cmd/pgterminate/main.go b/cmd/pgterminate/main.go
index 6c1c4ad..25cbcbf 100644
--- a/cmd/pgterminate/main.go
+++ b/cmd/pgterminate/main.go
@@ -47,6 +47,7 @@ func main() {
 	flag.StringVar(&config.IncludeUsersRegex, "include-users-regex", "", "Terminate users matching this regexp")
 	flag.Var(&config.ExcludeUsers, "exclude-user", "Ignore this user (can be called multiple times)")
 	flag.StringVar(&config.ExcludeUsersRegex, "exclude-users-regex", "", "Ignore users matching this regexp")
+	flag.BoolVar(&config.Cancel, "cancel", false, "Cancel sessions instead of terminate")
 	flag.Parse()
 
 	log.SetLevel(log.WarnLevel)
diff --git a/terminator/terminator.go b/terminator/terminator.go
index 7119145..b97d5c4 100644
--- a/terminator/terminator.go
+++ b/terminator/terminator.go
@@ -55,7 +55,11 @@ func (t *Terminator) Run() {
 
 // terminateAndNotify terminates a list of sessions and notifies channel
 func (t *Terminator) terminateAndNotify(sessions []base.Session) {
-	t.db.TerminateSessions(sessions)
+	if t.config.Cancel {
+		t.db.CancelSessions(sessions)
+	} else {
+		t.db.TerminateSessions(sessions)
+	}
 	for _, session := range sessions {
 		t.sessions <- session
 	}