diff --git a/README.md b/README.md index 18d7f8f..2f1b76e 100644 --- a/README.md +++ b/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. diff --git a/base/config.go b/base/config.go index c6e3778..71f2921 100644 --- a/base/config.go +++ b/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 +} diff --git a/base/utils.go b/base/utils.go index c7e7c8d..0962721 100644 --- a/base/utils.go +++ b/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 +} diff --git a/cmd/pgterminate/main.go b/cmd/pgterminate/main.go index 0c84b7f..6c1c4ad 100644 --- a/cmd/pgterminate/main.go +++ b/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) diff --git a/config.yaml.example b/config.yaml.example index 022357e..b3c2b2e 100644 --- a/config.yaml.example +++ b/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)" diff --git a/terminator/terminator.go b/terminator/terminator.go index 5749600..7119145 100644 --- a/terminator/terminator.go +++ b/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")