From 070566136c1ac72c4fe68df71732d0e917eca3a0 Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Fri, 8 Jan 2021 11:06:53 -0500
Subject: [PATCH] database: expand field map for composite fields

This allows models to fetch multiple SQL columns to obtain the data
necessary to compute a single composite GraphQL field, while still
avoiding unnecessary fetches if those fields are not queried.
---
 database/ql.go | 17 +++++++++++++----
 database/sq.go | 39 +++++++++++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/database/ql.go b/database/ql.go
index bb8fddad1f14ff95cc7c182259530929c4888f45..58e6fdf155e10d8c3463582974c01ce00e5d326b 100644
--- a/database/ql.go
+++ b/database/ql.go
@@ -3,6 +3,7 @@ package database
 import (
 	"context"
 	"sort"
+	"strings"
 
 	"github.com/lib/pq"
 	"github.com/vektah/gqlparser/v2/ast"
@@ -46,8 +47,10 @@ func Scan(ctx context.Context, m Model) []interface{} {
 
 	var fields []interface{}
 	for _, qlField := range qlFields {
-		if field, ok := m.Fields().GQL(qlField.Name); ok {
-			fields = append(fields, field.Ptr)
+		if gqlFields, ok := m.Fields().GQL(qlField.Name); ok {
+			for _, field := range gqlFields {
+				fields = append(fields, field.Ptr)
+			}
 		}
 	}
 
@@ -75,8 +78,14 @@ func Columns(ctx context.Context, m Model) []string {
 
 	var columns []string
 	for _, gql := range fields {
-		if field, ok := m.Fields().GQL(gql.Name); ok {
-			columns = append(columns, WithAlias(m.Alias(), field.SQL))
+		if sqlFields, ok := m.Fields().GQL(gql.Name); ok {
+			for _, field := range sqlFields {
+				sql := field.SQL
+				if !strings.ContainsRune(sql, '.') {
+					sql = WithAlias(m.Alias(), field.SQL)
+				}
+				columns = append(columns, sql)
+			}
 		}
 	}
 
diff --git a/database/sq.go b/database/sq.go
index cda4614933bbb5f175ae977a0597e2a36647afd2..fac913b0c3968769b99749b0603c1f8ca14ad4bc 100644
--- a/database/sq.go
+++ b/database/sq.go
@@ -19,8 +19,8 @@ type FieldMap struct {
 type ModelFields struct {
 	Fields []*FieldMap
 
-	byGQL map[string]*FieldMap
-	bySQL map[string]*FieldMap
+	byGQL map[string][]*FieldMap
+	bySQL map[string][]*FieldMap
 	anon  []*FieldMap
 }
 
@@ -29,19 +29,25 @@ func (mf *ModelFields) buildCache() {
 		return
 	}
 
-	mf.byGQL = make(map[string]*FieldMap)
-	mf.bySQL = make(map[string]*FieldMap)
+	mf.byGQL = make(map[string][]*FieldMap)
+	mf.bySQL = make(map[string][]*FieldMap)
 	for _, f := range mf.Fields {
 		if f.GQL != "" {
-			mf.byGQL[f.GQL] = f
+			if _, ok := mf.byGQL[f.GQL]; !ok {
+				mf.byGQL[f.GQL] = nil
+			}
+			mf.byGQL[f.GQL] = append(mf.byGQL[f.GQL], f)
 		} else {
 			mf.anon = append(mf.anon, f)
 		}
-		mf.bySQL[f.SQL] = f
+		if _, ok := mf.bySQL[f.SQL]; !ok {
+			mf.bySQL[f.SQL] = nil
+		}
+		mf.bySQL[f.SQL] = append(mf.bySQL[f.SQL], f)
 	}
 }
 
-func (mf *ModelFields) GQL(name string) (*FieldMap, bool) {
+func (mf *ModelFields) GQL(name string) ([]*FieldMap, bool) {
 	mf.buildCache()
 	if f, ok := mf.byGQL[name]; !ok {
 		return nil, false
@@ -50,7 +56,7 @@ func (mf *ModelFields) GQL(name string) (*FieldMap, bool) {
 	}
 }
 
-func (mf *ModelFields) SQL(name string) (*FieldMap, bool) {
+func (mf *ModelFields) SQL(name string) ([]*FieldMap, bool) {
 	mf.buildCache()
 	if f, ok := mf.bySQL[name]; !ok {
 		return nil, false
@@ -74,6 +80,11 @@ type Model interface {
 	Table() string
 }
 
+type ExtendedModel interface {
+	Model
+	Select(q sq.SelectBuilder) sq.SelectBuilder
+}
+
 func Select(ctx context.Context, cols ...interface{}) sq.SelectBuilder {
 	q := sq.Select().PlaceholderFormat(sq.Dollar)
 	for _, col := range cols {
@@ -83,7 +94,11 @@ func Select(ctx context.Context, cols ...interface{}) sq.SelectBuilder {
 		case []string:
 			q = q.Columns(col...)
 		case Model:
-			q = q.Columns(Columns(ctx, col)...)
+			if em, ok := col.(ExtendedModel); ok {
+				q = em.Select(q.Columns(Columns(ctx, col)...))
+			} else {
+				q = q.Columns(Columns(ctx, col)...)
+			}
 		default:
 			panic(fmt.Errorf("Unknown selectable type %T", col))
 		}
@@ -112,10 +127,14 @@ func Apply(m Model, input map[string]interface{}) sq.UpdateBuilder {
 	}()
 
 	for field, value := range input {
-		f, ok := m.Fields().GQL(field)
+		fields, ok := m.Fields().GQL(field)
 		if !ok {
 			continue
 		}
+		if len(fields) != 1 {
+			panic(fmt.Errorf("Apply cannot be used with composite fields"))
+		}
+		f := fields[0]
 
 		var (
 			pv reflect.Value = reflect.Indirect(reflect.ValueOf(f.Ptr))