mgocrud/lookup.go

package mgocrud

import (
	"fmt"
	"reflect"

	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// Lookup extends results with data for inline structs
func Lookup(db *mgo.Database, structField string, results interface{}, selector bson.M) error {

	t := reflect.TypeOf(results)
	v := reflect.ValueOf(results)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
		v = v.Elem()
	}

	isSlice := false
	elIsPtr := false
	if t.Kind() == reflect.Slice {
		t = t.Elem()
		if t.Kind() == reflect.Ptr {
			t = t.Elem()
			elIsPtr = true
		}
		isSlice = true
	}

	if t.Kind() != reflect.Struct {
		return fmt.Errorf("%+v is not a struct", t)
	}

	field, found := t.FieldByName(structField)
	if !found {
		return fmt.Errorf("struct field %s not found in struct %+v", structField, t)
	}
	fieldType := field.Type
	fieldIsPtr := false
	if fieldType.Kind() == reflect.Ptr {
		fieldType = fieldType.Elem()
		fieldIsPtr = true
	}

	fieldID, found := t.FieldByName(structField + "ID")
	if !found {
		return fmt.Errorf("struct field %s not found in struct %+v", structField+"ID", t)
	}
	idType := fieldID.Type
	idIsPtr := false
	if idType.Kind() == reflect.Ptr {
		idType = idType.Elem()
		idIsPtr = true
	}

	if idType != reflect.TypeOf(bson.ObjectId("")) {
		return fmt.Errorf("field %s in struct %+v is not bson.ObjectId", structField+"ID", t)
	}

	objectIDs := make(map[bson.ObjectId]interface{})

	// make slice for single elements to use generic code
	if !isSlice {
		tmpSlice := []interface{}{
			results,
		}
		v = reflect.ValueOf(tmpSlice)
		if reflect.TypeOf(results).Kind() == reflect.Ptr {
			elIsPtr = true
		}
	}

	l := v.Len()
	for i := 0; i < l; i++ {
		elV := v.Index(i)
		if elIsPtr {
			elV = elV.Elem()
		}
		if elV.Kind() == reflect.Ptr {
			// needed if tmpSlice -> dont know why double ptr resolv
			elV = elV.Elem()
		}
		lookupID := elV.FieldByName(fieldID.Name)
		if idIsPtr {
			lookupID = lookupID.Elem()
		}
		objectIDs[lookupID.Interface().(bson.ObjectId)] = nil
	}

	lArr := len(objectIDs)
	if lArr <= 0 {
		// no entries to map
		return nil
	}
	sArr := make([]bson.M, lArr, lArr)
	aI := 0
	for sID := range objectIDs {
		sArr[aI] = bson.M{
			"_id": sID,
		}
		aI++
	}
	sQuery := bson.M{
		"$or": sArr,
	}

	slice := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
	x := reflect.New(slice.Type())
	x.Elem().Set(slice)
	objectResults := x.Interface()
	objectSlice := x.Elem()

	objectIDIsPtr := false
	objectIDField, found := fieldType.FieldByName("ID")
	objectIDType := objectIDField.Type
	if found && objectIDType.Kind() == reflect.Ptr {
		objectIDType = objectIDType.Elem()
		objectIDIsPtr = true
	}
	if !found || objectIDType != reflect.TypeOf(bson.ObjectId("")) {
		return fmt.Errorf("ID type in objects struct %+v is not bson.ObjectId", fieldType)
	}

	err := ReadCollection(db, objectResults, sQuery, selector, 0, 0, nil, nil)
	if err != nil {
		panic(err)
	}

	// map IDs to object for better resolving
	oLen := objectSlice.Len()
	for i := 0; i < oLen; i++ {
		oID := objectSlice.Index(i).FieldByName("ID")
		if objectIDIsPtr {
			oID = oID.Elem()
		}
		objectIDs[oID.Interface().(bson.ObjectId)] = objectSlice.Index(i).Addr().Interface()
	}

	for i := 0; i < l; i++ {
		elV := v.Index(i)
		if elIsPtr {
			elV = elV.Elem()
		}
		if elV.Kind() == reflect.Ptr {
			// needed if tmpSlice -> dont know why double ptr resolv
			elV = elV.Elem()
		}

		oID := elV.FieldByName(structField + "ID")
		if idIsPtr {
			oID = oID.Elem()
		}
		objectID := oID.Interface().(bson.ObjectId)
		object := objectIDs[objectID]

		field := elV.FieldByName(structField)
		if fieldIsPtr {
			field.Set(reflect.ValueOf(object))
		} else {
			field.Set(reflect.ValueOf(object).Elem())
		}

	}

	return nil

}