go-rel/rel

package rel

import (
    "reflect"
    "strings"
    "sync"
    "time"

    "github.com/jinzhu/inflection"
    "github.com/serenize/snaker"
)

var (
    primariesCache    sync.Map
    documentMetaCache sync.Map
    rtTime            = reflect.TypeOf(time.Time{})
    rtBool            = reflect.TypeOf(false)
    rtInt             = reflect.TypeOf(int(0))
    rtTable           = reflect.TypeOf((*table)(nil)).Elem()
    rtPrimary         = reflect.TypeOf((*primary)(nil)).Elem()
)

// DocumentFlag stores information about document as a flag.
type DocumentFlag int8

// Is returns true if it's defined.
func (df DocumentFlag) Is(flag DocumentFlag) bool {
    return (df & flag) == flag
}

const (
    // Invalid flag.
    Invalid DocumentFlag = 1 << iota
    // HasCreatedAt flag.
    HasCreatedAt
    // HasUpdatedAt flag.
    HasUpdatedAt
    // HasDeletedAt flag.
    HasDeletedAt
    // HasDeleted flag.
    HasDeleted
    // Versioning
    HasVersioning
)

type table interface {
    Table() string
}

type primary interface {
    PrimaryFields() []string
    PrimaryValues() []any
}

type primaryData struct {
    field []string
    index [][]int
}

type cachedDocumentMeta struct {
    table        string
    index        map[string][]int
    fields       []string
    belongsTo    []string
    hasOne       []string
    hasMany      []string
    primaryField []string
    primaryIndex [][]int
    preload      []string
    flag         DocumentFlag
}

// Adds a prefix to field names
func appendWithPrefix(target, fieldNames []string, prefix string) []string {
    if prefix == "" {
        return append(target, fieldNames...)
    }
    for _, name := range fieldNames {
        target = append(target, prefix+name)
    }
    return target
}

// Adds a field index and checks for conflicts
func (cdm *cachedDocumentMeta) addFieldIndex(name string, index []int) {
    if _, ok := cdm.index[name]; ok {
        panic("rel: conflicting field (" + name + ") in struct")
    }
    cdm.index[name] = index
}

// Transfer values from other document data
func (cdm *cachedDocumentMeta) mergeEmbedded(other cachedDocumentMeta, indexPrefix int, namePrefix string) {
    for name, path := range other.index {
        cdm.addFieldIndex(namePrefix+name, append([]int{indexPrefix}, path...))
    }
    cdm.fields = appendWithPrefix(cdm.fields, other.fields, namePrefix)
    cdm.belongsTo = appendWithPrefix(cdm.belongsTo, other.belongsTo, namePrefix)
    cdm.hasOne = appendWithPrefix(cdm.hasOne, other.hasOne, namePrefix)
    cdm.hasMany = appendWithPrefix(cdm.hasMany, other.hasMany, namePrefix)
    cdm.primaryField = appendWithPrefix(cdm.primaryField, other.primaryField, namePrefix)
    for _, index := range other.primaryIndex {
        cdm.primaryIndex = append(cdm.primaryIndex, append([]int{indexPrefix}, index...))
    }
    cdm.preload = appendWithPrefix(cdm.preload, other.preload, namePrefix)
    cdm.flag |= other.flag
}

type DocumentMeta struct {
    rt reflect.Type
    cachedDocumentMeta
}

// Table returns name of the table.
func (dm DocumentMeta) Table() string {
    return dm.table
}

// PrimaryFields column name of this document.
func (dm DocumentMeta) PrimaryFields() []string {
    if len(dm.primaryField) == 0 {
        panic("rel: failed to infer primary key for type " + dm.rt.String())
    }

    return dm.primaryField
}

// PrimaryField column name of this document.
// panic if document uses composite key.
func (dm DocumentMeta) PrimaryField() string {
    if fields := dm.PrimaryFields(); len(fields) == 1 {
        return fields[0]
    }

    panic("rel: composite primary key is not supported")
}

// Index returns map of column name and it's struct index.
func (dm DocumentMeta) Index() map[string][]int {
    return dm.index
}

// Fields returns list of fields available on this document.
func (dm DocumentMeta) Fields() []string {
    return dm.fields
}

// Type returns reflect.Type of given field. if field does not exist, second returns value will be false.
func (dm DocumentMeta) Type(field string) (reflect.Type, bool) {
    if i, ok := dm.index[field]; ok {
        var (
            ft = dm.rt.FieldByIndex(i).Type
        )

        if ft.Kind() == reflect.Ptr {
            ft = ft.Elem()
        } else if ft.Kind() == reflect.Slice && ft.Elem().Kind() == reflect.Ptr {
            ft = reflect.SliceOf(ft.Elem().Elem())
        }

        return ft, true
    }

    return nil, false
}

// BelongsTo fields of this document.
func (dm DocumentMeta) BelongsTo() []string {
    return dm.belongsTo
}

// HasOne fields of this document.
func (dm DocumentMeta) HasOne() []string {
    return dm.hasOne
}

// HasMany fields of this document.
func (dm DocumentMeta) HasMany() []string {
    return dm.hasMany
}

// Preload fields of this document.
func (dm DocumentMeta) Preload() []string {
    return dm.preload
}

// Association of this document with given name.
func (dm DocumentMeta) Association(name string) AssociationMeta {
    if assoc, ok := dm.association(name); ok {
        return assoc
    }

    panic("rel: no field named (" + name + ") in type " + dm.rt.String() + " found ")
}

func (dm DocumentMeta) association(name string) (AssociationMeta, bool) {
    index, ok := dm.index[name]
    if !ok {
        return AssociationMeta{}, false
    }

    return getAssociationMeta(dm.rt, index), true
}

// Flag returns true if struct contains specified flag.
func (dm DocumentMeta) Flag(flag DocumentFlag) bool {
    return dm.flag.Is(flag)
}

func getDocumentMeta(rt reflect.Type, skipAssoc bool) DocumentMeta {
    if meta, cached := documentMetaCache.Load(rt); cached {
        return DocumentMeta{
            cachedDocumentMeta: meta.(cachedDocumentMeta),
            rt:                 rt,
        }
    }

    var (
        meta = cachedDocumentMeta{
            table: tableName(rt),
            index: make(map[string][]int, rt.NumField()),
        }
    )

    // TODO probably better to use slice index instead.
    for i := 0; i < rt.NumField(); i++ {
        var (
            sf           = rt.Field(i)
            typ          = sf.Type
            name, tagged = fieldName(sf)
        )

        if c := sf.Name[0]; c < 'A' || c > 'Z' || name == "" {
            continue
        }

        for typ.Kind() == reflect.Ptr || typ.Kind() == reflect.Interface || typ.Kind() == reflect.Slice {
            typ = typ.Elem()
        }

        if typ.Kind() == reflect.Struct && isEmbedded(sf) {
            embedded := getDocumentMeta(typ, skipAssoc)
            embeddedName := ""
            if tagged {
                embeddedName = name
            }
            meta.mergeEmbedded(embedded.cachedDocumentMeta, i, embeddedName)
            continue
        }

        meta.addFieldIndex(name, sf.Index)

        if flag := extractFlag(typ, name); flag != Invalid {
            meta.fields = append(meta.fields, name)
            meta.flag |= flag
            continue
        }

        if typ.Kind() != reflect.Struct {
            meta.fields = append(meta.fields, name)
            continue
        }

        // struct without primary key is a field
        // TODO: test by scanner/valuer instead?
        if pk, _ := searchPrimary(typ); len(pk) == 0 {
            meta.fields = append(meta.fields, name)
            continue
        }

        if !skipAssoc {
            var (
                assocMeta = getAssociationMeta(rt, sf.Index)
            )

            switch assocMeta.typ {
            case BelongsTo:
                meta.belongsTo = append(meta.belongsTo, name)
            case HasOne:
                meta.hasOne = append(meta.hasOne, name)
            case HasMany:
                meta.hasMany = append(meta.hasMany, name)
            }

            if assocMeta.autoload {
                meta.preload = append(meta.preload, name)
            }
        }
    }

    primaryField, primaryIndex := searchPrimary(rt)
    meta.primaryField = append(meta.primaryField, primaryField...)
    meta.primaryIndex = append(meta.primaryIndex, primaryIndex...)

    if !skipAssoc {
        documentMetaCache.Store(rt, meta)
    }

    return DocumentMeta{
        rt:                 rt,
        cachedDocumentMeta: meta,
    }
}

func extractTimeFlag(name string) DocumentFlag {
    switch name {
    case "created_at", "inserted_at":
        return HasCreatedAt
    case "updated_at":
        return HasUpdatedAt
    case "deleted_at":
        return HasDeletedAt
    }
    return Invalid
}

func extractBoolFlag(name string) DocumentFlag {
    if name == "deleted" {
        return HasDeleted
    }
    return Invalid
}

func extractIntFlag(name string) DocumentFlag {
    if name == "lock_version" {
        return HasVersioning
    }
    return Invalid
}

func extractFlag(rt reflect.Type, name string) DocumentFlag {
    if rt == rtTime {
        return extractTimeFlag(name)
    }
    if rt == rtBool {
        return extractBoolFlag(name)
    }
    if rt == rtInt {
        return extractIntFlag(name)
    }
    return Invalid
}

func fieldName(sf reflect.StructField) (string, bool) {
    if tag := sf.Tag.Get("db"); tag != "" {
        name := strings.Split(tag, ",")[0]

        if name == "-" {
            return "", true
        }

        if name != "" {
            return name, true
        }
    }

    return snaker.CamelToSnake(sf.Name), false
}

func isEmbedded(sf reflect.StructField) bool {
    // anonymous structs are always embedded
    if sf.Anonymous {
        return true
    }
    if tag := sf.Tag.Get("db"); strings.HasSuffix(tag, ",embedded") {
        return true
    }
    return false
}

func searchPrimary(rt reflect.Type) ([]string, [][]int) {
    if result, cached := primariesCache.Load(rt); cached {
        p := result.(primaryData)
        return p.field, p.index
    }

    var (
        field         []string
        index         [][]int
        fallbackIndex = -1
    )

    if rt.Implements(rtPrimary) {
        var (
            v = reflect.Zero(rt).Interface().(primary)
        )

        field = v.PrimaryFields()
        // index kept nil to mark interface usage
    } else {
        for i := 0; i < rt.NumField(); i++ {
            sf := rt.Field(i)

            if tag := sf.Tag.Get("db"); strings.HasSuffix(tag, ",primary") {
                index = append(index, sf.Index)
                name, _ := fieldName(sf)
                field = append(field, name)
                continue
            }

            // check fallback for id field
            if strings.EqualFold("id", sf.Name) {
                fallbackIndex = i
            }
        }
    }

    if len(field) == 0 && fallbackIndex >= 0 {
        field = []string{"id"}
        index = [][]int{{fallbackIndex}}
    }

    primariesCache.Store(rt, primaryData{
        field: field,
        index: index,
    })

    return field, index
}

func tableName(rt reflect.Type) string {
    var name string
    if rt.Implements(rtTable) {
        name = reflect.Zero(rt).Interface().(table).Table()
    } else {
        name = inflection.Plural(rt.Name())
        name = snaker.CamelToSnake(name)
    }

    return name
}