feat: Include and exclude databases
- Add filters to include and exclude strings - Use filters to include and exclude sessions (user and databases supported) - Add tests to filters and terminator Signed-off-by: Julien Riou <julien@riou.xyz>
This commit is contained in:
parent
29dbbc5bef
commit
8709ee542b
9 changed files with 553 additions and 53 deletions
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
144
terminator/terminator_test.go
Normal file
144
terminator/terminator_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue