diff --git a/Makefile b/Makefile index bae228a..6261d94 100644 --- a/Makefile +++ b/Makefile @@ -15,5 +15,9 @@ release: go build -ldflags "${LDFLAGS}" -o bin/${BINARY} cmd/${BINARY}/main.go (cd bin && tar czf ${BINARY}-${APPVERSION}-${GOOS}-${GOARCH}.tar.gz ${BINARY}) +test: + go test base/* + go test terminator/* + clean: rm -rf bin diff --git a/README.md b/README.md index 920670e..429cf18 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,10 @@ When include users list or regex is set, `pgterminate` will focus on included us When exclude users list or regex is set and no include option is set, `pgterminate` will terminate all sessions except excluded users. +# Filtering databases + +Similar to users, `pgterminate` is able to filter sessions by database. The same set of options are available if you replace `user` by `database` (ex: `-include-user` to `-include-database` and so on). + # Listeners LISTEN queries are asynchronous. Sessions are set to "idle" state even if they are waiting for messages to be sent to the queue. `pgterminate` can exclude sessions in that state by looking at the last known query starting with "LISTEN", with the `exclude-listeners` parameter. diff --git a/base/config.go b/base/config.go index 2099652..f9ff9c3 100644 --- a/base/config.go +++ b/base/config.go @@ -17,31 +17,41 @@ 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"` - LogFormat string `yaml:"log-format"` - 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 - ExcludeListeners bool `yaml:"exclude-listeners"` - Cancel bool `yaml:"cancel"` + 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"` + LogFormat string `yaml:"log-format"` + 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 + IncludeUsersFilters []Filter + ExcludeUsers StringFlags `yaml:"exclude-users"` + ExcludeUsersRegex string `yaml:"exclude-users-regex"` + ExcludeUsersRegexCompiled *regexp.Regexp + ExcludeUsersFilters []Filter + IncludeDatabases StringFlags `yaml:"include-databases"` + IncludeDatabasesRegex string `yaml:"include-databases-regex"` + IncludeDatabasesRegexCompiled *regexp.Regexp + IncludeDatabasesFilters []Filter + ExcludeDatabases StringFlags `yaml:"exclude-databases"` + ExcludeDatabasesRegex string `yaml:"exclude-databases-regex"` + ExcludeDatabasesRegexCompiled *regexp.Regexp + ExcludeDatabasesFilters []Filter + ExcludeListeners bool `yaml:"exclude-listeners"` + Cancel bool `yaml:"cancel"` } func init() { @@ -83,6 +93,7 @@ func (c *Config) Reload() { } err := c.CompileRegexes() Panic(err) + c.CompileFilters() } // Dsn formats a connection string based on Config @@ -129,6 +140,43 @@ func (c *Config) CompileRegexes() (err error) { return nil } +// CompileFilters creates Filter objects based on patterns and compiled regexp +func (c *Config) CompileFilters() { + + c.IncludeUsersFilters = nil + if c.IncludeUsers != nil { + c.IncludeUsersFilters = append(c.IncludeUsersFilters, NewIncludeFilter(c.IncludeUsers)) + } + if c.IncludeUsersRegexCompiled != nil { + c.IncludeUsersFilters = append(c.IncludeUsersFilters, NewIncludeFilterRegex(c.IncludeUsersRegexCompiled)) + } + + c.ExcludeUsersFilters = nil + if c.ExcludeUsers != nil { + c.ExcludeUsersFilters = append(c.ExcludeUsersFilters, NewExcludeFilter(c.ExcludeUsers)) + } + if c.ExcludeUsersRegexCompiled != nil { + c.ExcludeUsersFilters = append(c.ExcludeUsersFilters, NewExcludeFilterRegex(c.ExcludeUsersRegexCompiled)) + } + + c.IncludeDatabasesFilters = nil + if c.IncludeDatabases != nil { + c.IncludeDatabasesFilters = append(c.IncludeDatabasesFilters, NewIncludeFilter(c.IncludeDatabases)) + } + if c.IncludeDatabasesRegexCompiled != nil { + c.IncludeDatabasesFilters = append(c.IncludeDatabasesFilters, NewIncludeFilterRegex(c.IncludeDatabasesRegexCompiled)) + } + + c.ExcludeDatabasesFilters = nil + if c.ExcludeDatabases != nil { + c.ExcludeDatabasesFilters = append(c.ExcludeDatabasesFilters, NewExcludeFilter(c.ExcludeDatabases)) + } + if c.ExcludeDatabasesRegexCompiled != nil { + c.ExcludeDatabasesFilters = append(c.ExcludeDatabasesFilters, NewExcludeFilterRegex(c.ExcludeDatabasesRegexCompiled)) + } + +} + // StringFlags append multiple string flags into a string slice type StringFlags []string diff --git a/base/filter.go b/base/filter.go new file mode 100644 index 0000000..7798acd --- /dev/null +++ b/base/filter.go @@ -0,0 +1,122 @@ +package base + +import ( + "fmt" + "reflect" + "regexp" +) + +// Filter interface to tell if a string should be included or not +type Filter interface { + Include(string) bool + String() string +} + +// IncludeFilter to include a string when it's included in a list of strings +type IncludeFilter struct { + patterns []string +} + +// NewIncludeFilter to create an IncludeFilter +func NewIncludeFilter(patterns []string) IncludeFilter { + return IncludeFilter{ + patterns: patterns, + } +} + +// Include returns true when a string is included in a list of patterns +// Implements the Filter interface +func (f IncludeFilter) Include(s string) bool { + // No or empty patterns must include + if f.patterns == nil || reflect.DeepEqual(f.patterns, []string{""}) { + return true + } + return InSlice(s, f.patterns) +} + +// String to pretty print an IncludeFilter +// Implements the Filter interface +func (f IncludeFilter) String() string { + return fmt.Sprintf("", f.patterns) +} + +// IncludeFilterRegex to include a string when it matches a regex +type IncludeFilterRegex struct { + regex *regexp.Regexp +} + +// NewIncludeFilterRegex to create an IncludeFilterRegex +func NewIncludeFilterRegex(regex *regexp.Regexp) IncludeFilterRegex { + return IncludeFilterRegex{ + regex: regex, + } +} + +// Include returns true when the string matches the regex +// Implements the Filter interface +func (f IncludeFilterRegex) Include(s string) bool { + if f.regex == nil || f.regex.MatchString(s) { + return true + } + return false +} + +// String to pretty print an IncludeFilterRegex +// Implements the Filter interface +func (f IncludeFilterRegex) String() string { + return fmt.Sprintf("", f.regex.String()) +} + +// ExcludeFilter to include a string when it's not included in a list of strings +type ExcludeFilter struct { + patterns []string +} + +// NewExcludeFilter to create an ExcludeFilter +func NewExcludeFilter(patterns []string) ExcludeFilter { + return ExcludeFilter{ + patterns: patterns, + } +} + +// Include returns true when the string is not included in the patterns +// Implements the Filter interface +func (f ExcludeFilter) Include(s string) bool { + return !InSlice(s, f.patterns) +} + +// String to pretty print an ExcludeFilter +// Implements the Filter interface +func (f ExcludeFilter) String() string { + return fmt.Sprintf("", f.patterns) +} + +// ExcludeFilterRegex to include a string when it doesnn't match a regex +type ExcludeFilterRegex struct { + regex *regexp.Regexp +} + +// NewExcludeFilterRegex to create an ExcludeFilterRegex +func NewExcludeFilterRegex(regex *regexp.Regexp) ExcludeFilterRegex { + return ExcludeFilterRegex{ + regex: regex, + } +} + +// Include returns true when the string doesn't match the regex +// Implements the Filter interface +func (f ExcludeFilterRegex) Include(s string) bool { + if f.regex == nil || f.regex.MatchString("") { + return true + } + if f.regex.MatchString(s) { + return false + } + return true +} + +// String to pretty print an ExcludeFilterRegex +// Implements the Filter interface +func (f ExcludeFilterRegex) String() string { + return fmt.Sprintf("", f.regex.String()) +} diff --git a/base/filter_test.go b/base/filter_test.go new file mode 100644 index 0000000..bf75eb5 --- /dev/null +++ b/base/filter_test.go @@ -0,0 +1,124 @@ +package base + +import ( + "fmt" + "regexp" + "testing" +) + +func TestIncludeFilter(t *testing.T) { + tests := []struct { + name string + value string + patterns []string + wanted bool + }{ + {"No filter", "test", nil, true}, + {"Empty filter", "test", []string{""}, true}, + {"Single pattern matching", "test", []string{"test"}, true}, + {"Multiple patterns matching", "test", []string{"test", "postgres"}, true}, + {"Single pattern with no match", "nomatch", []string{"test"}, false}, + {"Multiple patterns with no match", "nomatch", []string{"test", "postgres"}, false}, + } + + for _, tc := range tests { + t.Run(fmt.Sprintf(tc.name), func(t *testing.T) { + f := NewIncludeFilter(tc.patterns) + + if got := f.Include(tc.value); got != tc.wanted { + t.Errorf("Included must be %t for patterns '%s'", tc.wanted, tc.patterns) + } else { + t.Logf("Included is %t for patterns '%s'", tc.wanted, tc.patterns) + } + }) + } +} + +func TestIncludeFilterRegex(t *testing.T) { + tests := []struct { + name string + value string + regex string + wanted bool + }{ + {"No filter", "test", "", true}, + {"String pattern matching", "test", "test", true}, + {"Regex patterns matching", "test", "^t(.*)$", true}, + {"String pattern with no match", "nomatch", "test", false}, + {"Regex patterns with no match", "nomatch", "^t(.*)$", false}, + } + + for _, tc := range tests { + t.Run(fmt.Sprintf(tc.name), func(t *testing.T) { + compiledRegex, err := regexp.Compile(tc.regex) + if err != nil { + t.Fatalf("Regex '%s' doesn't compile: %v", tc.regex, err) + } + + f := NewIncludeFilterRegex(compiledRegex) + if got := f.Include(tc.value); got != tc.wanted { + t.Errorf("Included must be %t for regex '%s'", tc.wanted, tc.regex) + } else { + t.Logf("Included is %t for regex '%s'", tc.wanted, tc.regex) + } + }) + } +} + +func TestExcludeFilter(t *testing.T) { + tests := []struct { + name string + value string + patterns []string + wanted bool + }{ + {"No filter", "test", nil, true}, + {"Empty filter", "test", []string{""}, true}, + {"Single pattern matching", "test", []string{"test"}, false}, + {"Multiple patterns matching", "test", []string{"test", "postgres"}, false}, + {"Single pattern with no match", "nomatch", []string{"test"}, true}, + {"Multiple patterns with no match", "nomatch", []string{"test", "postgres"}, true}, + } + + for _, tc := range tests { + t.Run(fmt.Sprintf(tc.name), func(t *testing.T) { + f := NewExcludeFilter(tc.patterns) + if got := f.Include(tc.value); got != tc.wanted { + t.Errorf("Included must be %t for patterns '%s'", tc.wanted, tc.patterns) + } else { + t.Logf("Included is %t for patterns '%s'", tc.wanted, tc.patterns) + } + }) + } +} + +func TestExcludeFilterRegex(t *testing.T) { + tests := []struct { + name string + value string + regex string + wanted bool + }{ + {"No filter", "test", "", true}, + {"String pattern matching", "test", "test", false}, + {"Regex patterns matching", "test", "^t(.*)$", false}, + {"String pattern with no match", "nomatch", "test", true}, + {"Regex patterns with no match", "nomatch", "^t(.*)$", true}, + } + + for _, tc := range tests { + t.Run(fmt.Sprintf(tc.name), func(t *testing.T) { + compiledRegex, err := regexp.Compile(tc.regex) + if err != nil { + t.Fatalf("Regex '%s' doesn't compile: %v", tc.regex, err) + } + + f := NewExcludeFilterRegex(compiledRegex) + if got := f.Include(tc.value); got != tc.wanted { + t.Errorf("Included must be %t for regex '%s'", tc.wanted, tc.regex) + } else { + t.Logf("Included is %t for regex '%s'", tc.wanted, tc.regex) + } + }) + } +} diff --git a/cmd/pgterminate/main.go b/cmd/pgterminate/main.go index cecc413..e88be55 100644 --- a/cmd/pgterminate/main.go +++ b/cmd/pgterminate/main.go @@ -58,6 +58,10 @@ 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.Var(&config.IncludeDatabases, "include-database", "Terminate only this database (can be called multiple times)") + flag.StringVar(&config.IncludeDatabasesRegex, "include-databases-regex", "", "Terminate databases matching this regexp") + flag.Var(&config.ExcludeDatabases, "exclude-database", "Ignore this database (can be called multiple times)") + flag.StringVar(&config.ExcludeDatabasesRegex, "exclude-databases-regex", "", "Ignore databases matching this regexp") flag.BoolVar(&config.ExcludeListeners, "exclude-listeners", false, "Ignore sessions listening for events") flag.BoolVar(&config.Cancel, "cancel", false, "Cancel sessions instead of terminate") flag.Parse() @@ -112,6 +116,7 @@ func main() { err = config.CompileRegexes() base.Panic(err) + config.CompileFilters() if config.PidFile != "" { writePid(config.PidFile) diff --git a/config.yaml.example b/config.yaml.example index fb97dc1..d041bb9 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -21,4 +21,12 @@ # user1 # user2 #exclude-users-regex: "(user1|user2)" +#include-databases: +# db1 +# db2 +#include-databases-regex: "(db1|db2)" +#exclude-databases: +# db1 +# db2 +#exclude-databases-regex: "(db1|db2)" #cancel \ No newline at end of file diff --git a/terminator/terminator.go b/terminator/terminator.go index f5a6c87..f6b6f57 100644 --- a/terminator/terminator.go +++ b/terminator/terminator.go @@ -1,10 +1,11 @@ package terminator import ( - "github.com/jouir/pgterminate/base" - "github.com/jouir/pgterminate/log" "strings" "time" + + "github.com/jouir/pgterminate/base" + "github.com/jouir/pgterminate/log" ) // Terminator looks for sessions, filters actives and idles, terminate them and notify sessions channel @@ -71,30 +72,6 @@ func (t *Terminator) notify(sessions []*base.Session) { } } -// filterUsers 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) filterUsers(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 -} - // filterListeners excludes sessions with last query starting with "LISTEN" func (t *Terminator) filterListeners(sessions []*base.Session) (filtered []*base.Session) { for _, session := range sessions { @@ -105,11 +82,75 @@ func (t *Terminator) filterListeners(sessions []*base.Session) (filtered []*base return filtered } +// filterUsers include and exclude users based on filters +func (t *Terminator) filterUsers(sessions []*base.Session) []*base.Session { + + var included []*base.Session + if t.config.IncludeUsersFilters == nil { + included = sessions + } else { + for _, filter := range t.config.IncludeUsersFilters { + for _, session := range sessions { + if filter.Include(session.User) { + included = append(included, session) + } + } + } + } + + var filtered []*base.Session + if t.config.ExcludeUsersFilters == nil { + filtered = included + } else { + for _, filter := range t.config.ExcludeUsersFilters { + for _, session := range included { + if filter.Include(session.User) { + filtered = append(filtered, session) + } + } + } + } + + return filtered +} + +// filterUsers include and exclude databases based on filters +func (t *Terminator) filterDatabases(sessions []*base.Session) []*base.Session { + + var included []*base.Session + if t.config.IncludeDatabasesFilters == nil { + included = sessions + } else { + for _, filter := range t.config.IncludeDatabasesFilters { + for _, session := range sessions { + if filter.Include(session.Db) { + included = append(included, session) + } + } + } + } + + var filtered []*base.Session + if t.config.ExcludeDatabasesFilters == nil { + filtered = included + } else { + for _, filter := range t.config.ExcludeDatabasesFilters { + for _, session := range included { + if filter.Include(session.Db) { + filtered = append(filtered, session) + } + } + } + } + + return filtered +} + // filter executes all filter functions on a list of sessions func (t *Terminator) filter(sessions []*base.Session) (filtered []*base.Session) { - filtered = sessions + filtered = t.filterListeners(sessions) filtered = t.filterUsers(filtered) - filtered = t.filterListeners(filtered) + filtered = t.filterDatabases(filtered) return filtered } diff --git a/terminator/terminator_test.go b/terminator/terminator_test.go new file mode 100644 index 0000000..f68b528 --- /dev/null +++ b/terminator/terminator_test.go @@ -0,0 +1,144 @@ +package terminator + +import ( + "reflect" + "testing" + + "github.com/jouir/pgterminate/base" +) + +func TestFilterUsers(t *testing.T) { + + sessions := []*base.Session{ + {User: "test"}, + {User: "test_1"}, + {User: "test_2"}, + {User: "postgres"}, + } + + tests := []struct { + name string + config *base.Config + want []*base.Session + }{ + { + "No filter", + &base.Config{}, + sessions, + }, + { + "Include a single user", + &base.Config{IncludeUsers: []string{"test"}}, + []*base.Session{{User: "test"}}, + }, + { + "Include multiple users", + &base.Config{IncludeUsers: []string{"test_1", "test_2"}}, + []*base.Session{{User: "test_1"}, {User: "test_2"}}, + }, + { + "Exclude a single user", + &base.Config{ExcludeUsers: []string{"test"}}, + []*base.Session{{User: "test_1"}, {User: "test_2"}, {User: "postgres"}}, + }, + { + "Exclude multiple users", + &base.Config{ExcludeUsers: []string{"test_1", "test_2"}}, + []*base.Session{{User: "test"}, {User: "postgres"}}, + }, + { + "Include multiple users and exclude one", + &base.Config{IncludeUsers: []string{"test", "test_1", "test_2"}, ExcludeUsers: []string{"test"}}, + []*base.Session{{User: "test_1"}, {User: "test_2"}}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tc.config.CompileFilters() + terminator := &Terminator{config: tc.config} + got := terminator.filterUsers(sessions) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("got %+v; want %+v", ListUsers(got), ListUsers(tc.want)) + } else { + t.Logf("Success") + } + }) + } +} + +// ListUsers extract usernames from a list of sessions +func ListUsers(sessions []*base.Session) (users []string) { + for _, session := range sessions { + users = append(users, session.User) + } + return users +} + +func TestFilterDatabases(t *testing.T) { + + sessions := []*base.Session{ + {Db: "test"}, + {Db: "test_1"}, + {Db: "test_2"}, + {Db: "postgres"}, + } + + tests := []struct { + name string + config *base.Config + want []*base.Session + }{ + { + "No filter", + &base.Config{}, + sessions, + }, + { + "Include a single database", + &base.Config{IncludeDatabases: []string{"test"}}, + []*base.Session{{Db: "test"}}, + }, + { + "Include multiple databases", + &base.Config{IncludeDatabases: []string{"test_1", "test_2"}}, + []*base.Session{{Db: "test_1"}, {Db: "test_2"}}, + }, + { + "Exclude a single database", + &base.Config{ExcludeDatabases: []string{"test"}}, + []*base.Session{{Db: "test_1"}, {Db: "test_2"}, {Db: "postgres"}}, + }, + { + "Exclude multiple databases", + &base.Config{ExcludeDatabases: []string{"test_1", "test_2"}}, + []*base.Session{{Db: "test"}, {Db: "postgres"}}, + }, + { + "Include multiple databases and exclude one", + &base.Config{IncludeDatabases: []string{"test", "test_1", "test_2"}, ExcludeDatabases: []string{"test"}}, + []*base.Session{{Db: "test_1"}, {Db: "test_2"}}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tc.config.CompileFilters() + terminator := &Terminator{config: tc.config} + got := terminator.filterDatabases(sessions) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("got %+v; want %+v", ListDatabases(got), ListDatabases(tc.want)) + } else { + t.Logf("Success") + } + }) + } +} + +// ListDatabases extract usernames from a list of sessions +func ListDatabases(sessions []*base.Session) (databases []string) { + for _, session := range sessions { + databases = append(databases, session.Db) + } + return databases +}