Add Auto Role Feature

cmd/goff/main.go

@@ -51,6 +51,7 @@ func main() {
 
 	us := &postgres.UserService{DB: postgres.DB}
 	gs := &postgres.GuildService{DB: postgres.DB}
+	rs := &postgres.RoleService{DB: postgres.DB}
 
 	//prefixes := []string{
 	//	"Go.",
@@ -71,7 +72,7 @@ func main() {
 
 	// Add Command Handlers
 	exts.AddCommandHandlers(&manager)
-	services.InitalizeServices(us, gs)
+	services.InitalizeServices(us, gs, rs)
 
 	//if _, ok := handler.Commands["help"]; !ok {
 	//	handler.AddDefaultHelpCommand()

internal/exts/roles/roles.go

@@ -7,7 +7,9 @@ import (
 
 	"github.com/bwmarrin/discordgo"
 	"github.com/dustinpianalto/disgoman"
+	"github.com/dustinpianalto/goff"
 	"github.com/dustinpianalto/goff/internal/postgres"
+	"github.com/dustinpianalto/goff/internal/services"
 )
 
 var MakeRoleSelfAssignableCommand = &disgoman.Command{
@@ -254,3 +256,134 @@ func memberHasRole(m *discordgo.Member, id string) bool {
 	}
 	return false
 }
+
+var MakeAutoRoleCommand = &disgoman.Command{
+	Name:                "make-auto-role",
+	Aliases:             []string{"makear"},
+	Description:         "Marks the passed in role to auto add to new members.",
+	OwnerOnly:           false,
+	Hidden:              false,
+	RequiredPermissions: disgoman.PermissionManageServer,
+	Invoke:              makeAutoRoleCommandFunc,
+}
+
+func makeAutoRoleCommandFunc(ctx disgoman.Context, args []string) {
+	var roleString string
+	var roleID string
+	if len(args) > 0 {
+		roleString = strings.Join(args, " ")
+		if strings.HasPrefix(roleString, "<@&") && strings.HasSuffix(roleString, ">") {
+			roleID = roleString[3 : len(roleString)-1]
+		} else if _, err := strconv.Atoi(roleString); err == nil {
+			roleID = roleString
+		} else {
+			for _, role := range ctx.Guild.Roles {
+				if roleString == role.Name {
+					roleID = role.ID
+				}
+			}
+		}
+	}
+	fmt.Println(roleID)
+	var r *discordgo.Role
+	var err error
+	if r, err = ctx.Session.State.Role(ctx.Guild.ID, roleID); err != nil {
+		ctx.CommandManager.ErrorChannel <- disgoman.CommandError{
+			Context: ctx,
+			Message: "Can't find that Role.",
+			Error:   err,
+		}
+		return
+	}
+	role, err := services.RoleService.Role(r.ID)
+	if err != nil {
+		role = &goff.Role{
+			ID:             r.ID,
+			IsModerator:    false,
+			IsAdmin:        false,
+			SelfAssignable: false,
+			AutoRole:       true,
+			Guild:          ctx.Guild.ID,
+		}
+		role, err = services.RoleService.AddRole(role)
+		if err != nil {
+			ctx.CommandManager.ErrorChannel <- disgoman.CommandError{
+				Context: ctx,
+				Message: "Error making an Auto Role",
+				Error:   err,
+			}
+			return
+		}
+	} else {
+		err = services.RoleService.MakeAutoRole(role)
+		if err != nil {
+			ctx.CommandManager.ErrorChannel <- disgoman.CommandError{
+				Context: ctx,
+				Message: "Error making an Auto Role",
+				Error:   err,
+			}
+			return
+		}
+	}
+
+}
+
+var RemoveAutoRoleCommand = &disgoman.Command{
+	Name:                "remove-auto-role",
+	Aliases:             []string{"removear"},
+	Description:         "Remove role from those added to new members.",
+	OwnerOnly:           false,
+	Hidden:              false,
+	RequiredPermissions: disgoman.PermissionManageServer,
+	Invoke:              removeAutoRoleCommandFunc,
+}
+
+func removeAutoRoleCommandFunc(ctx disgoman.Context, args []string) {
+	var roleString string
+	var roleID string
+	if len(args) > 0 {
+		roleString = strings.Join(args, " ")
+		if strings.HasPrefix(roleString, "<@&") && strings.HasSuffix(roleString, ">") {
+			roleID = roleString[3 : len(roleString)-1]
+		} else if _, err := strconv.Atoi(roleString); err == nil {
+			roleID = roleString
+		} else {
+			for _, role := range ctx.Guild.Roles {
+				if roleString == role.Name {
+					roleID = role.ID
+				}
+			}
+		}
+	}
+	fmt.Println(roleID)
+	var r *discordgo.Role
+	var err error
+	if r, err = ctx.Session.State.Role(ctx.Guild.ID, roleID); err != nil {
+		ctx.CommandManager.ErrorChannel <- disgoman.CommandError{
+			Context: ctx,
+			Message: "Can't find that Role.",
+			Error:   err,
+		}
+		return
+	}
+	role, err := services.RoleService.Role(r.ID)
+	if err != nil {
+		ctx.CommandManager.ErrorChannel <- disgoman.CommandError{
+			Context: ctx,
+			Message: "Role was not configured as an Auto Role",
+			Error:   err,
+		}
+		return
+	} else {
+		err = services.RoleService.RemoveAutoRole(role)
+		if err != nil {
+			ctx.CommandManager.ErrorChannel <- disgoman.CommandError{
+				Context: ctx,
+				Message: "Error removing Auto Role",
+				Error:   err,
+			}
+			return
+		}
+	}
+
+}

internal/exts/user_management/member_join.go

@@ -41,6 +41,22 @@ func OnGuildMemberAdd(s *discordgo.Session, member *discordgo.GuildMemberAdd) {
 	if err != nil {
 		log.Println("Error adding user to guild: ", err)
 	}
+	roles, err := services.RoleService.GetAutoRoles(member.GuildID)
+	if err != nil {
+		log.Println("Error getting Auto Join Roles: ", err)
+	}
+	for _, r := range roles {
+		role, err := s.State.Role(member.GuildID, r.ID)
+		if err != nil {
+			log.Println("Error getting role: ", err)
+			continue
+		}
+		err = s.GuildMemberRoleAdd(member.GuildID, member.User.ID, role.ID)
+		if err != nil {
+			log.Println("Error adding Role to member: ", err)
+			continue
+		}
+	}
 }
 
 func OnGuildMemberRemove(s *discordgo.Session, member *discordgo.GuildMemberRemove) {

internal/postgres/postfixes.go

@@ -32,6 +32,10 @@ var postfixes = []postfix{
 		Name:   "4_Add_Multi_Column_Unique_XUsersGuilds",
 		Invoke: addMultiColumnUniqueXUsersGuilds,
 	},
+	postfix{
+		Name:   "5_Add_Auto_Role_To_Roles_Table",
+		Invoke: addAutoRoleToRoles,
+	},
 }
 
 func RunPostfixes() {
@@ -162,3 +166,20 @@ func addMultiColumnUniqueXUsersGuilds(revert bool) error {
 	}
 	return nil
 }
+
+func addAutoRoleToRoles(revert bool) error {
+	var queryString string
+	if !revert {
+		queryString = `ALTER TABLE roles
+						ADD COLUMN auto_role bool DEFAULT false`
+	} else {
+		queryString = `ALTER TABLE roles
+						DROP COLUMN auto_role`
+	}
+	_, err := DB.Exec(queryString)
+	if err != nil {
+		log.Println(err)
+		return err
+	}
+	return nil
+}

internal/postgres/roles.go

@@ -0,0 +1,86 @@
+package postgres
+
+import (
+	"database/sql"
+	"log"
+
+	"github.com/dustinpianalto/goff"
+)
+
+type RoleService struct {
+	DB *sql.DB
+}
+
+func (s *RoleService) Role(id string) (*goff.Role, error) {
+	var r goff.Role
+	queryString := `SELECT id, moderator, admin, self_assignable, guild_id FROM roles WHERE id = $1`
+	row := s.DB.QueryRow(queryString, id)
+	if err := row.Scan(&r.ID, &r.IsModerator, &r.IsAdmin, &r.SelfAssignable, &r.Guild); err != nil {
+		return nil, err
+	}
+	return &r, nil
+}
+
+func (s *RoleService) AddRole(r *goff.Role) (*goff.Role, error) {
+	queryString := `INSERT INTO roles (id, moderator, admin, self_assignable, auto_role, guild)
+						VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT DO NOTHING`
+	_, err := s.DB.Exec(queryString, r.ID, r.IsModerator, r.IsAdmin, r.SelfAssignable, r.AutoRole, r.Guild)
+	if err != nil {
+		return nil, err
+	}
+	return r, nil
+}
+
+func (s *RoleService) DeleteRole(r *goff.Role) error {
+	queryString := `DELETE FROM roles WHERE id = $1`
+	_, err := s.DB.Exec(queryString, r.ID)
+	return err
+}
+
+func (s *RoleService) MakeSelfAssignable(r *goff.Role) error {
+	queryString := `UPDATE roles SET self_assignable = true WHERE id = $1`
+	_, err := s.DB.Exec(queryString, r.ID)
+	return err
+}
+
+func (s *RoleService) RemoveSelfAssignable(r *goff.Role) error {
+	queryString := `UPDATE roles SET self_assignable = false WHERE id = $1`
+	_, err := s.DB.Exec(queryString, r.ID)
+	return err
+}
+
+func (s *RoleService) MakeAutoRole(r *goff.Role) error {
+	queryString := `UPDATE roles SET auto_role = true WHERE id = $1`
+	_, err := s.DB.Exec(queryString, r.ID)
+	return err
+}
+
+func (s *RoleService) RemoveAutoRole(r *goff.Role) error {
+	queryString := `UPDATE roles SET auto_role = false WHERE id = $1`
+	_, err := s.DB.Exec(queryString, r.ID)
+	return err
+}
+
+func (s *RoleService) GetAutoRoles(gID string) ([]*goff.Role, error) {
+	var roles []*goff.Role
+	queryString := `SELECT id FROM roles WHERE guild_id = $1`
+	rows, err := s.DB.Query(queryString, gID)
+	if err != nil {
+		return nil, err
+	}
+	for rows.Next() {
+		var id string
+		err := rows.Scan(&id)
+		if err != nil {
+			log.Println(err)
+			continue
+		}
+		role, err := s.Role(id)
+		if err != nil {
+			log.Println(err)
+			continue
+		}
+		roles = append(roles, role)
+	}
+	return roles, nil
+}

internal/services/services.go

@@ -4,8 +4,10 @@ import "github.com/dustinpianalto/goff"
 
 var UserService goff.UserService
 var GuildService goff.GuildService
+var RoleService goff.RoleService
 
-func InitalizeServices(us goff.UserService, gs goff.GuildService) {
+func InitalizeServices(us goff.UserService, gs goff.GuildService, rs goff.RoleService) {
 	UserService = us
 	GuildService = gs
+	RoleService = rs
 }

role.go

@@ -0,0 +1,21 @@
+package goff
+
+type Role struct {
+	ID             string
+	IsModerator    bool
+	IsAdmin        bool
+	SelfAssignable bool
+	AutoRole       bool
+	Guild          string
+}
+
+type RoleService interface {
+	Role(id string) (*Role, error)
+	AddRole(r *Role) (*Role, error)
+	DeleteRole(r *Role) error
+	MakeSelfAssignable(r *Role) error
+	RemoveSelfAssignable(r *Role) error
+	MakeAutoRole(r *Role) error
+	RemoveAutoRole(r *Role) error
+	GetAutoRoles(gID string) ([]*Role, error)
+}