Add cancel option to terminate queries

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.

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

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

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)

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
 	}