Merge 3d3c403 into be3a613

cmd/cli/config.go

@@ -8,6 +8,13 @@ type config struct {
 	LogFile  string          `yaml:"log_file"`
 	LogLevel string          `yaml:"log_level"`
 	Redis    cmd.RedisConfig `yaml:"redis"`
+	Cleanup  cleanupConfig   `yaml:"cleanup"`
+}
+
+type cleanupConfig struct {
+	Whitelist       []string `yaml:"whitelist"`
+	Delete          bool     `yaml:"delete"`
+	HandleAnonymous bool     `yaml:"handle_anonymous"`
 }
 
 func getDefault() config {
@@ -19,5 +26,8 @@ func getDefault() config {
 			Port:            "6379",
 			ConnectionLimit: 512,
 		},
+		Cleanup: cleanupConfig{
+			Whitelist: []string{},
+		},
 	}
 }

cmd/cli/main.go

@@ -42,13 +42,14 @@ var (
 )
 
 var (
+	cleanup  = flag.Bool("cleanup", false, "User Existence Check")
 	userDel  = flag.String("user-del", "", "Delete all contacts and subscriptions for a user")
 	fromUser = flag.String("from-user", "", "Transfer subscriptions and contacts from user.")
 	toUser   = flag.String("to-user", "", "Transfer subscriptions and contacts to user.")
 )
 
 func main() {
-	logger, dataBase := initApp()
+	confCleanup, logger, dataBase := initApp()
 
 	if *update {
 		fromVersion := checkValidVersion(logger, updateFromVersion, true)
@@ -77,6 +78,7 @@ func main() {
 			logger.Errorf("Failed to enable images in all notifications")
 		}
 	}
+
 	if *fromUser != "" || *toUser != "" {
 		if err := transferUserSubscriptionsAndContacts(dataBase, *fromUser, *toUser); err != nil {
 			logger.Error(err)
@@ -88,9 +90,16 @@ func main() {
 			logger.Error(err)
 		}
 	}
+
+	if *cleanup {
+		logger.Debugf("User whitelist: %#v", confCleanup.Whitelist)
+		if err := handleCleanup(logger, dataBase, confCleanup); err != nil {
+			logger.Error(err)
+		}
+	}
 }
 
-func initApp() (moira.Logger, moira.Database) {
+func initApp() (cleanupConfig, moira.Logger, moira.Database) {
 	flag.Parse()
 	if *printVersion {
 		fmt.Println("Moira - alerting system based on graphite data")
@@ -120,7 +129,7 @@ func initApp() (moira.Logger, moira.Database) {
 
 	databaseSettings := config.Redis.GetSettings()
 	dataBase := redis.NewDatabase(logger, databaseSettings, redis.Cli)
-	return logger, dataBase
+	return config.Cleanup, logger, dataBase
 }
 
 func checkValidVersion(logger moira.Logger, updateFromVersion *string, isUpdate bool) string {

cmd/cli/user.go

@@ -1,6 +1,9 @@
 package main
 
 import (
+	"encoding/json"
+	"os"
+
 	"github.com/moira-alert/moira"
 )
 
@@ -75,3 +78,93 @@ func deleteUser(database moira.Database, user string) error {
 
 	return nil
 }
+
+func handleCleanup(logger moira.Logger, database moira.Database, config cleanupConfig) error {
+	var users []string
+
+	reader := json.NewDecoder(os.Stdin)
+
+	if err := reader.Decode(&users); err != nil {
+		return err
+	}
+
+	return usersCleanup(logger, database, users, config)
+}
+
+func usersCleanup(logger moira.Logger, database moira.Database, users []string, config cleanupConfig) error {
+	if config.HandleAnonymous {
+		config.Whitelist = append(config.Whitelist, "")
+	}
+
+	usersMap := make(map[string]bool, len(users)+len(config.Whitelist))
+
+	for _, user := range append(users, config.Whitelist...) {
+		usersMap[user] = true
+	}
+
+	contacts, err := database.GetAllContacts()
+	if err != nil {
+		return err
+	}
+
+	if len(contacts) == 0 {
+		return nil
+	}
+
+	usersNotFound := make(map[string]bool, len(contacts))
+
+	for _, contact := range contacts {
+		if contact == nil {
+			continue
+		}
+
+		if !usersMap[contact.User] {
+			usersNotFound[contact.User] = true
+		}
+	}
+
+	for user := range usersNotFound {
+		if config.Delete {
+			if err = deleteUser(database, user); err != nil {
+				break
+			}
+		} else {
+			if err = offNotification(database, user); err != nil {
+				break
+			}
+		}
+
+		logger.Debug("User cleaned:", user)
+	}
+
+	return err
+}
+
+func offNotification(database moira.Database, user string) error {
+	subscriptionIDs, err := database.GetUserSubscriptionIDs(user)
+	if err != nil {
+		return err
+	}
+
+	subscriptions, err := database.GetSubscriptions(subscriptionIDs)
+	if err != nil {
+		return err
+	}
+
+	saveSubscriptions := make([]*moira.SubscriptionData, 0, len(subscriptions))
+
+	for _, subscription := range subscriptions {
+		if subscription == nil {
+			continue
+		}
+
+		if !subscription.Enabled {
+			continue
+		}
+
+		subscription.Enabled = false
+		saveSubscriptions = append(saveSubscriptions, subscription)
+	}
+
+	return database.SaveSubscriptions(saveSubscriptions)
+}

cmd/cli/user_test.go

@@ -0,0 +1,124 @@
+package main
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/moira-alert/moira/database/redis"
+	"github.com/moira-alert/moira/logging/go-logging"
+
+	"github.com/moira-alert/moira"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+func TestUpdateUsers(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode.")
+	}
+
+	conf := getDefault()
+	logger, err := logging.ConfigureLog(conf.LogFile, conf.LogLevel, "cli")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	database := redis.NewDatabase(logger, conf.Redis.GetSettings(), redis.Cli)
+	conf.Cleanup.Whitelist = []string{"Nikolay", ""}
+
+	users := []string{"Aleksey", "Arkadiy", "Emil"}
+
+	Convey("Test clean users in moira", t, func() {
+		if err := createDefaultData(database); err != nil {
+			t.Fatal(err)
+		}
+
+		defer func(t *testing.T) {
+			if err := cleanData(database); err != nil {
+				t.Fatal(err)
+			}
+		}(t)
+
+		Convey("Test off notifications", func() {
+			So(usersCleanup(logger, database, users, conf.Cleanup), ShouldBeNil)
+			for _, contact := range contacts {
+				subscription, err := database.GetSubscription("subscription_" + contact.ID)
+
+				So(err, ShouldBeNil)
+
+				if strings.Contains(subscription.User, "Another") {
+					So(subscription.Enabled, ShouldBeFalse)
+				} else {
+					So(subscription.Enabled, ShouldBeTrue)
+				}
+			}
+		})
+
+		Convey("Verify deletion of contacts and subscriptions", func() {
+			conf.Cleanup.Delete = true
+			So(usersCleanup(logger, database, users, conf.Cleanup), ShouldBeNil)
+			for _, contact := range contacts {
+				if !strings.Contains(contact.User, "Another") {
+					continue
+				}
+
+				_, err := database.GetSubscription("subscription_" + contact.ID)
+				So(err, ShouldNotBeNil)
+
+				_, err = database.GetContact(contact.ID)
+				So(err, ShouldNotBeNil)
+			}
+		})
+	})
+
+}
+
+func createDefaultData(database moira.Database) error {
+	subscriptions := make([]*moira.SubscriptionData, 0, len(contacts))
+
+	for _, contact := range contacts {
+		if err := database.SaveContact(contact); err != nil {
+			return err
+		}
+
+		subscriptions = append(subscriptions,
+			&moira.SubscriptionData{ID: "subscription_" + contact.ID,
+				User:     contact.User,
+				Enabled:  true,
+				Tags:     []string{"Tag" + contact.User},
+				Contacts: []string{contact.ID},
+			},
+		)
+	}
+
+	if err := database.SaveSubscriptions(subscriptions); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func cleanData(database moira.Database) error {
+
+	for _, contact := range contacts {
+		if err := database.RemoveContact(contact.ID); err != nil {
+			return err
+		}
+
+		if err := database.RemoveSubscription("subscription_" + contact.ID); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+var contacts = []*moira.ContactData{
+	{ID: "cli_id_00000000000001", User: "Aleksey"},
+	{ID: "cli_id_00000000000002", User: "Arkadiy"},
+	{ID: "cli_id_00000000000003", User: "Emil"},
+	{ID: "cli_id_00000000000004", User: "Nikolay"},
+	{ID: "cli_id_00000000000005", User: "Another1"},
+	{ID: "cli_id_00000000000006", User: "Another2"},
+	{ID: "cli_id_00000000000007", User: ""},
+}

database/redis/subscription.go

@@ -296,6 +296,9 @@ func (connector *DbConnector) getSubscriptionsTriggers(subscriptions []*moira.Su
 			return triggers, err
 		}
 		for _, trigger := range subscriptionTriggers {
+			if trigger == nil {
+				continue
+			}
 			triggersMap[trigger.ID] = trigger
 		}
 	}