Add users filtering

README.md

@@ -47,5 +47,47 @@ Print usage:
 pgterminate -help
 ```
 
+# Filtering users
+
+`pgterminate` is able to include or exclude users from being terminated.
+
+## Configuration
+### List
+Arguments `-include-user` or `-exclude-user` can be used multiple times for multiple users:
+
+```
+pgterminate -include-user user1 -include-user user2
+```
+Or in configuration file:
+
+```
+include-users:
+  user1
+  user2
+```
+Same applies for `-exclude-user` (argument) and `exclude-users` (file).
+
+### Regexes
+Regexes can be configured:
+
+```
+pgterminate -include-users-regex "(user1|user2)"
+```
+Or in configuration file:
+
+```
+include-users-regex: "(user1|user2)"
+```
+
+Same applies for `-exclude-users-regex` (argument) and `exclude-users-regex` (file).
+
+## Include users
+
+When include users list or regex is set, `pgterminate` will focus on included users only. It could terminate excluded users if any. If you want to exclude users, use exclude options only.
+
+## Exclude users
+
+When exclude users list or regex is set and no include option is set, `pgterminate` will terminate all sessions except excluded users.
+
 # License
 `pgterminate` is released under [The Unlicense](https://github.com/jouir/pgterminate/blob/master/LICENSE) license. Code is under public domain.

base/config.go

@@ -6,6 +6,7 @@ import (
 	"gopkg.in/yaml.v2"
 	"io/ioutil"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"sync"
 )
@@ -15,22 +16,28 @@ var AppName string
 
 // Config receives configuration options
 type Config struct {
-	mutex          sync.Mutex
-	File           string
-	Host           string  `yaml:"host"`
-	Port           int     `yaml:"port"`
-	User           string  `yaml:"user"`
-	Password       string  `yaml:"password"`
-	Database       string  `yaml:"database"`
-	Interval       float64 `yaml:"interval"`
-	ConnectTimeout int     `yaml:"connect-timeout"`
-	IdleTimeout    float64 `yaml:"idle-timeout"`
-	ActiveTimeout  float64 `yaml:"active-timeout"`
-	LogDestination string  `yaml:"log-destination"`
-	LogFile        string  `yaml:"log-file"`
-	PidFile        string  `yaml:"pid-file"`
-	SyslogIdent    string  `yaml:"syslog-ident"`
-	SyslogFacility string  `yaml:"syslog-facility"`
+	mutex                     sync.Mutex
+	File                      string
+	Host                      string      `yaml:"host"`
+	Port                      int         `yaml:"port"`
+	User                      string      `yaml:"user"`
+	Password                  string      `yaml:"password"`
+	Database                  string      `yaml:"database"`
+	Interval                  float64     `yaml:"interval"`
+	ConnectTimeout            int         `yaml:"connect-timeout"`
+	IdleTimeout               float64     `yaml:"idle-timeout"`
+	ActiveTimeout             float64     `yaml:"active-timeout"`
+	LogDestination            string      `yaml:"log-destination"`
+	LogFile                   string      `yaml:"log-file"`
+	PidFile                   string      `yaml:"pid-file"`
+	SyslogIdent               string      `yaml:"syslog-ident"`
+	SyslogFacility            string      `yaml:"syslog-facility"`
+	IncludeUsers              StringFlags `yaml:"include-users"`
+	IncludeUsersRegex         string      `yaml:"include-users-regex"`
+	IncludeUsersRegexCompiled *regexp.Regexp
+	ExcludeUsers              StringFlags `yaml:"exclude-users"`
+	ExcludeUsersRegex         string      `yaml:"exclude-users-regex"`
+	ExcludeUsersRegexCompiled *regexp.Regexp
 }
 
 func init() {
@@ -62,7 +69,7 @@ func (c *Config) Read(file string) error {
 	return nil
 }
 
-// Reload reads from file and update configuration
+// Reload reads from file to update configuration and re-compile regexes
 func (c *Config) Reload() {
 	log.Debug("Reloading configuration")
 	c.mutex.Lock()
@@ -70,6 +77,8 @@ func (c *Config) Reload() {
 	if c.File != "" {
 		c.Read(c.File)
 	}
+	err := c.CompileRegexes()
+	Panic(err)
 }
 
 // Dsn formats a connection string based on Config
@@ -98,3 +107,34 @@ func (c *Config) Dsn() string {
 	}
 	return strings.Join(parameters, " ")
 }
+
+// CompileRegexes transforms regexes from string to regexp instance
+func (c *Config) CompileRegexes() (err error) {
+	if c.IncludeUsersRegex != "" {
+		c.IncludeUsersRegexCompiled, err = regexp.Compile(c.IncludeUsersRegex)
+		if err != nil {
+			return err
+		}
+	}
+	if c.ExcludeUsersRegex != "" {
+		c.ExcludeUsersRegexCompiled, err = regexp.Compile(c.ExcludeUsersRegex)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// StringFlags append multiple string flags into a string slice
+type StringFlags []string
+
+// String for implementing flag interface
+func (s *StringFlags) String() string {
+	return "multiple strings flag"
+}
+
+// Set adds alues into the slice
+func (s *StringFlags) Set(value string) error {
+	*s = append(*s, value)
+	return nil
+}

base/utils.go

@@ -10,3 +10,13 @@ func Panic(err error) {
 		log.Fatalf("%s\n", err)
 	}
 }
+
+// InSlice detects value presence in a string slice
+func InSlice(value string, slice []string) bool {
+	for _, val := range slice {
+		if value == val {
+			return true
+		}
+	}
+	return false
+}

cmd/pgterminate/main.go

@@ -43,6 +43,10 @@ func main() {
 	flag.StringVar(&config.PidFile, "pid-file", "", "Write process id into a file")
 	flag.StringVar(&config.SyslogIdent, "syslog-ident", "pgterminate", "Define syslog tag")
 	flag.StringVar(&config.SyslogFacility, "syslog-facility", "", "Define syslog facility from LOCAL0 to LOCAL7")
+	flag.Var(&config.IncludeUsers, "include-user", "Terminate only this user (can be called multiple times)")
+	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.Parse()
 
 	log.SetLevel(log.WarnLevel)
@@ -93,6 +97,9 @@ func main() {
 		}
 	}
 
+	err = config.CompileRegexes()
+	base.Panic(err)
+
 	if config.PidFile != "" {
 		writePid(config.PidFile)
 		defer removePid(config.PidFile)

config.yaml.example

@@ -12,3 +12,11 @@ pid-file: /var/run/pgterminate/pgterminate.pid
 #log-destination: console|file|syslog
 #syslog-ident: pgterminate
 #syslog-facility: LOCAL0
+#include-users:
+#    user1
+#    user2
+#include-users-regex: "(user1|user2)"
+#exclude-users:
+#    user1
+#    user2
+#exclude-users-regex: "(user1|user2)"

terminator/terminator.go

@@ -40,12 +40,12 @@ func (t *Terminator) Run() {
 			sessions := t.db.Sessions()
 			if t.config.ActiveTimeout != 0 {
 				actives := activeSessions(sessions, t.config.ActiveTimeout)
-				t.terminateAndNotify(actives)
+				t.terminateAndNotify(t.filter(actives))
 			}
 
 			if t.config.IdleTimeout != 0 {
 				idles := idleSessions(sessions, t.config.IdleTimeout)
-				t.terminateAndNotify(idles)
+				t.terminateAndNotify(t.filter(idles))
 			}
 			time.Sleep(time.Duration(t.config.Interval*1000) * time.Millisecond)
 		}
@@ -61,6 +61,30 @@ func (t *Terminator) terminateAndNotify(sessions []base.Session) {
 	}
 }
 
+// filter removes sessions according to include and exclude users settings
+// when include users slice and regex are not set, append all sessions except excluded users
+// otherwise, append included users
+func (t *Terminator) filter(sessions []base.Session) (filtered []base.Session) {
+	includeUsers, includeRegex := t.config.IncludeUsers, t.config.IncludeUsersRegexCompiled
+	excludeUsers, excludeRegex := t.config.ExcludeUsers, t.config.ExcludeUsersRegexCompiled
+
+	for _, session := range sessions {
+		if t.config.IncludeUsers == nil && includeRegex == nil {
+			// append all sessions except excluded users
+			if !base.InSlice(session.User, excludeUsers) || (excludeRegex != nil && !excludeRegex.MatchString(session.User)) {
+				filtered = append(filtered, session)
+			}
+		} else {
+			// append included users only
+			if base.InSlice(session.User, includeUsers) || (includeRegex != nil && includeRegex.MatchString(session.User)) {
+				filtered = append(filtered, session)
+			}
+		}
+	}
+
+	return filtered
+}
+
 // terminate terminates gracefully
 func (t *Terminator) terminate() {
 	log.Info("Disconnecting from instance")