Organization Tree Model in Mongo

admin_user_tree_model.go

// follow https://docs.mongodb.com/manual/tutorial/model-tree-structures-with-materialized-paths/
package models

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"qbox.us/mgo2"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

var (
	AdminUserTree *_AdminUserTree

	ColAdminUserTree = bson.M{
		"name": "admin_user_trees",
		"index": []string{
			"path",
		},
		"unique": []string{
			"name",
		},
	}
)

type (
	// AdminUserTreePath represents the path of a tree ndoe.
	AdminUserTreePath struct {
		Paths []bson.ObjectId
	}

	// AdminUserTreeModel represents a node of organization tree.
	AdminUserTreeModel struct {
		ID          bson.ObjectId      `bson:"_id"`
		Name        string             `bson:"name"`
		Path        *AdminUserTreePath `bson:"path"`
		ManagerID   bson.ObjectId      `bson:"manager_id,omitempty"`
		MemberIDs   []bson.ObjectId    `bson:"member_ids"`
		ViewerIDs   []bson.ObjectId    `bson:"viewer_ids"`
		HasChildren bool               `bson:"has_children"`
		SfProfileID string             `bson:"sf_profile_id,omitempty"`

		CreatedAt *time.Time `bson:"created_at"`
		UpdatedAt *time.Time `bson:"updated_at"`
	}

	_AdminUserTree struct{}
)

/**************************************/
/******** AdminUserTreePath ***********/
/**************************************/

// GetBSON implements bson.Getter interface.
func (m *AdminUserTreePath) GetBSON() (interface{}, error) {
	return m.String(), nil
}

// SetBSON implements bson.Setter interface.
func (m *AdminUserTreePath) SetBSON(raw bson.Raw) (err error) {
	var str string
	if err = raw.Unmarshal(str); err != nil {
		return err
	}
	m, err = parseAdminUserTreePath(str)
	return
}

func (m *AdminUserTreePath) String() (str string) {
	str = ","
	for _, p := range m.Paths {
		str += p.Hex() + ","
	}
	return
}

func parseAdminUserTreePath(str string) (path *AdminUserTreePath, err error) {
	if !strings.HasPrefix(str, ",") {
		err = errors.New(`path should start with ","`)
		return
	}

	str = strings.Trim(str, ",")
	ids := strings.Split(str, ",")
	path = &AdminUserTreePath{
		Paths: make([]bson.ObjectId, len(ids)),
	}
	for i, id := range ids {
		if !bson.IsObjectIdHex(id) {
			err = fmt.Errorf("bad object id string: %s", id)
			return
		}
		path.Paths[i] = bson.ObjectIdHex(id)
	}
	return
}

// RegEx returns the bson.RegEx with given options.
// Valid options as of this writing are 'i' for case insensitive matching,
// 'm' for multi-line matching, 'x' for verbose mode, 'l' to make \w, \W,
// and similar be locale-dependent, 's' for dot-all mode (a '.' matches everything),
// and 'u' to make \w, \W, and similar match unicode.
func (m *AdminUserTreePath) RegEx(options ...string) bson.RegEx {
	return bson.RegEx{
		Pattern: "/" + m.String() + "/",
		Options: strings.Join(options, ""),
	}
}

/**************************************/
/******** AdminUserTreeModel **********/
/**************************************/

func getAdminUerTreeCollection() *mgo.Collection {
	return mgo2.CopyCollection(DefaultAdminMongo.C(ColAdminUserTree))
}

// NewAdminUserTreeModel instances a *AdminUserTreeModel.
func NewAdminUserTreeModel() *AdminUserTreeModel {
	now := bson.Now()
	return &AdminUserTreeModel{
		ID:        bson.NewObjectId(),
		CreatedAt: &now,
	}
}

// pathAsParent returns the path when this AdminUserTreeModel as parent.
// e.g. path of this -> ",a,b,c,", id of this -> "d", pathAsParent of this -> ",a,b,c,d,"
func (m *AdminUserTreeModel) pathAsParent() string {
	prePath:=m.Path.String()
	if prePath==""{
		prePath=","
	}
	return prePath + m.ID.Hex() + ","
}

// Save upserts this instance.
func (m *AdminUserTreeModel) Save() error {
	if m.ID == "" {
		return errors.New("missing id")
	}

	c := getAdminUerClassCollection()
	defer mgo2.CloseCollection(c)

	n := bson.Now()
	m.UpdatedAt = &n

	_, err := c.UpsertId(m.ID, m)
	return err
}

// SetParent sets a AdminUserTreeModel as the parent.
// If the parent hasn't had a child before, congratulations!
// Mark the parent has already had a child and move the members to its child.
func (m *AdminUserTreeModel) SetParent(parent *AdminUserTreeModel) (err error) {
	c := getAdminUerTreeCollection()
	defer mgo2.CloseCollection(c)

	t := bson.Now()
	newPath := parent.pathAsParent()
	if _, err = c.UpdateAll(
		bson.M{"path": m.Path.RegEx()},
		bson.M{"$set": bson.M{
			"path":       newPath,
			"created_at": &t,
		}},
	); err != nil {
		return
	}

	if !parent.HasChildren {
		parent.HasChildren = true
		memberIDs := parent.MemberIDs
		parent.MemberIDs = nil
		if err = c.UpdateId(parent.ID, parent); err != nil {
			return
		}

		if memberIDs != nil && len(memberIDs) > 0 {
			if err = c.UpdateId(m.ID, bson.M{"$set": bson.M{
				"member_ids": memberIDs,
				"created_at": &t,
			}}); err != nil {
				return
			}
			m.MemberIDs = memberIDs
		}
	}

	m.Path, _ = parseAdminUserTreePath(newPath)
	return
}

// AddMembers adds members into this AdminUserTreeModel.
func (m *AdminUserTreeModel) AddMembers(memberIDs ...bson.ObjectId) {
	if m.MemberIDs == nil {
		m.MemberIDs = make([]bson.ObjectId, 0)
	}
	m.MemberIDs = append(m.MemberIDs, memberIDs...)
}