diff --git a/main.go b/main.go index 4e5b2f3..2ab3c81 100644 --- a/main.go +++ b/main.go @@ -88,6 +88,7 @@ func main() { // Migrate database schema db.AutoMigrate(&models.User{}) + db.AutoMigrate(&models.Alias{}) db.AutoMigrate(&models.Heartbeat{}).AddForeignKey("user_id", "users(id)", "RESTRICT", "RESTRICT") // Custom migrations and initial data @@ -95,9 +96,10 @@ func main() { migrateLanguages(db, config) // Services + aliasSrvc := &services.AliasService{config, db} heartbeatSrvc := &services.HeartbeatService{config, db} userSrvc := &services.UserService{config, db} - summarySrvc := &services.SummaryService{config, db, heartbeatSrvc} + summarySrvc := &services.SummaryService{config, db, heartbeatSrvc, aliasSrvc} // Handlers heartbeatHandler := &routes.HeartbeatHandler{HeartbeatSrvc: heartbeatSrvc} diff --git a/models/alias.go b/models/alias.go new file mode 100644 index 0000000..691ebdd --- /dev/null +++ b/models/alias.go @@ -0,0 +1,9 @@ +package models + +type Alias struct { + ID uint `gorm:"primary_key"` + Type uint8 `gorm:"not null; index:idx_alias_type_key"` + UserID string `gorm:"not null; index:idx_alias_user"` + Key string `gorm:"not null; index:idx_alias_type_key"` + Value string `gorm:"not null"` +} diff --git a/routes/summary.go b/routes/summary.go index 4e1ddf6..529159f 100644 --- a/routes/summary.go +++ b/routes/summary.go @@ -2,6 +2,7 @@ package routes import ( "crypto/md5" + "log" "net/http" "strconv" "time" @@ -44,6 +45,13 @@ func (h *SummaryHandler) Get(w http.ResponseWriter, r *http.Request) { } user := r.Context().Value(models.UserKey).(*models.User) + + // Initialize aliases for user + if !h.SummarySrvc.AliasService.IsInitialized(user.ID) { + log.Printf("Initializing aliases for user '%s'\n", user.ID) + h.SummarySrvc.AliasService.InitUser(user.ID) + } + params := r.URL.Query() interval := params.Get("interval") from, err := utils.ParseDate(params.Get("from")) diff --git a/services/alias.go b/services/alias.go new file mode 100644 index 0000000..dd92a3e --- /dev/null +++ b/services/alias.go @@ -0,0 +1,50 @@ +package services + +import ( + "errors" + + "github.com/jinzhu/gorm" + "github.com/n1try/wakapi/models" +) + +type AliasService struct { + Config *models.Config + Db *gorm.DB +} + +var userAliases map[string][]*models.Alias + +func (srv *AliasService) InitUser(userId string) error { + if userAliases == nil { + userAliases = make(map[string][]*models.Alias) + } + + var aliases []*models.Alias + if err := srv.Db. + Where(&models.Alias{UserID: userId}). + Find(&aliases).Error; err != nil { + return err + } + + userAliases[userId] = aliases + return nil +} + +func (srv *AliasService) GetAliasOrDefault(userId string, summaryType uint8, value string) (string, error) { + if userAliases, ok := userAliases[userId]; ok { + for _, a := range userAliases { + if a.Type == summaryType && a.Value == value { + return a.Key, nil + } + } + return value, nil + } + return "", errors.New("User aliases not initialized") +} + +func (src *AliasService) IsInitialized(userId string) bool { + if _, ok := userAliases[userId]; ok { + return true + } + return false +} diff --git a/services/summary.go b/services/summary.go index 084b1e3..caac8cc 100644 --- a/services/summary.go +++ b/services/summary.go @@ -12,6 +12,7 @@ type SummaryService struct { Config *models.Config Db *gorm.DB HeartbeatService *HeartbeatService + AliasService *AliasService } func (srv *SummaryService) GetSummary(from, to time.Time, user *models.User) (*models.Summary, error) { @@ -29,7 +30,7 @@ func (srv *SummaryService) GetSummary(from, to time.Time, user *models.User) (*m c := make(chan models.SummaryItemContainer) for _, t := range types { - go srv.aggregateBy(heartbeats, t, c) + go srv.aggregateBy(heartbeats, t, user, c) } for i := 0; i < len(types); i++ { @@ -60,7 +61,7 @@ func (srv *SummaryService) GetSummary(from, to time.Time, user *models.User) (*m return summary, nil } -func (srv *SummaryService) aggregateBy(heartbeats []*models.Heartbeat, summaryType uint8, c chan models.SummaryItemContainer) { +func (srv *SummaryService) aggregateBy(heartbeats []*models.Heartbeat, summaryType uint8, user *models.User, c chan models.SummaryItemContainer) { durations := make(map[string]time.Duration) for i, h := range heartbeats { @@ -80,6 +81,10 @@ func (srv *SummaryService) aggregateBy(heartbeats []*models.Heartbeat, summaryTy key = "unknown" } + if aliasedKey, err := srv.AliasService.GetAliasOrDefault(user.ID, summaryType, key); err == nil { + key = aliasedKey + } + if _, ok := durations[key]; !ok { durations[key] = time.Duration(0) } diff --git a/services/user.go b/services/user.go index 08809b1..3377032 100644 --- a/services/user.go +++ b/services/user.go @@ -5,8 +5,6 @@ import ( "github.com/n1try/wakapi/models" ) -const TableUser = "user" - type UserService struct { Config *models.Config Db *gorm.DB