gorm/callback_query.go

package gorm

import (
	"fmt"
	"reflect"

	"github.com/jinzhu/gorm"
)

func Query(scope *Scope) {
	defer scope.Trace(NowFunc())

	var (
		isSlice        bool
		isPtr          bool
		anyRecordFound bool
		destType       reflect.Type
		primaryKeys    []interface{}
	)

	var dest = scope.IndirectValue()
	if value, ok := scope.InstanceGet("gorm:query_destination"); ok {
		dest = reflect.Indirect(reflect.ValueOf(value))
	}

	if orderBy, ok := scope.InstanceGet("gorm:order_by_primary_key"); ok {
		if primaryKey := scope.PrimaryKey(); primaryKey != "" {
			scope.Search = scope.Search.clone().order(fmt.Sprintf("%v.%v %v", scope.QuotedTableName(), primaryKey, orderBy))
		}
	}

	if dest.Kind() == reflect.Slice {
		isSlice = true
		destType = dest.Type().Elem()
		if destType.Kind() == reflect.Ptr {
			isPtr = true
			destType = destType.Elem()
		}
	} else {
		scope.Search = scope.Search.clone().limit(1)
	}

	scope.prepareQuerySql()

	if !scope.HasError() {
		rows, err := scope.DB().Query(scope.Sql, scope.SqlVars...)
		scope.db.RowsAffected = 0

		if scope.Err(err) != nil {
			return
		}

		preloadMap := map[string]map[string]*gorm.Field{}

		columns, _ := rows.Columns()
		defer rows.Close()
		for rows.Next() {
			scope.db.RowsAffected += 1

			anyRecordFound = true
			elem := dest
			if isSlice {
				elem = reflect.New(destType).Elem()
			}

			var values = make([]interface{}, len(columns))

			fields := scope.New(elem.Addr().Interface()).Fields()
			var primaryKey interface{}
			for index, column := range columns {
				if field, ok := fields[column]; ok {
					if field.Field.Kind() == reflect.Ptr {
						values[index] = field.Field.Addr().Interface()
					} else {
						values[index] = reflect.New(reflect.PtrTo(field.Field.Type())).Interface()
					}
					if field.IsPrimaryKey {
						primaryKey = values[index]
						primaryKeys = append(primaryKeys, primaryKey)
					}
				} else {
					var value interface{}
					values[index] = &value
				}
			}

			scope.Err(rows.Scan(values...))

			for index, column := range columns {
				value := values[index]
				if field, ok := fields[column]; ok {
					if field.Field.Kind() == reflect.Ptr {
						field.Field.Set(reflect.ValueOf(value).Elem())
					} else if v := reflect.ValueOf(value).Elem().Elem(); v.IsValid() {
						field.Field.Set(v)
					}
				}
			}

			if isSlice {
				if isPtr {
					dest.Set(reflect.Append(dest, elem.Addr()))
				} else {
					dest.Set(reflect.Append(dest, elem))
				}
			}

			if scope.Search.Preload != nil {
				for key := range scope.Search.Preload {
					if field := fields[key]; field != nil {
						if preloadMap[key] == nil {
							preloadMap[key] = map[string]reflect.Value{}
						}
						preloadMap[key][fmt.Sprintf("%v", primaryKey)] = field
					}
				}
			}
		}

		for _, value := range preloadMap {
			var typ reflect.Type
			var relation *Relation
			for _, v := range value {
				typ = v.Field.Type()
				relation = v.Relationship
				break
			}
			sliceType := reflect.SliceOf(typ)
			slice := reflect.MakeSlice(sliceType, 0, 0)
			slicePtr := reflect.New(sliceType)
			slicePtr.Elem().Set(slice)
			if relation == "has_many" {
				scope.NewDB().Find(slicePtr.Interface(), primaryKeys)
			}
		}

		if !anyRecordFound && !isSlice {
			scope.Err(RecordNotFound)
		}
	}
}

func AfterQuery(scope *Scope) {
	scope.CallMethod("AfterFind")
}

func init() {
	DefaultCallback.Query().Register("gorm:query", Query)
	DefaultCallback.Query().Register("gorm:after_query", AfterQuery)
}