From 497551202238b0befd67d69b1da547bfff660948 Mon Sep 17 00:00:00 2001
From: Grace Noah <gracenoahgh@gmail.com>
Date: Sun, 11 Nov 2018 03:02:02 +0000
Subject: [PATCH] accept an optional ctx parameter on model methods

---
 codegen/object.go | 25 +++++++++++++++----------
 codegen/util.go   | 14 +++++++++++++-
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/codegen/object.go b/codegen/object.go
index c037344fe65..70a1e99ecbc 100644
--- a/codegen/object.go
+++ b/codegen/object.go
@@ -33,16 +33,17 @@ type Object struct {
 
 type Field struct {
 	*Type
-	Description    string          // Description of a field
-	GQLName        string          // The name of the field in graphql
-	GoFieldType    GoFieldType     // The field type in go, if any
-	GoReceiverName string          // The name of method & var receiver in go, if any
-	GoFieldName    string          // The name of the method or var in go, if any
-	Args           []FieldArgument // A list of arguments to be passed to this field
-	ForceResolver  bool            // Should be emit Resolver method
-	NoErr          bool            // If this is bound to a go method, does that method have an error as the second argument
-	Object         *Object         // A link back to the parent object
-	Default        interface{}     // The default value
+	Description      string          // Description of a field
+	GQLName          string          // The name of the field in graphql
+	GoFieldType      GoFieldType     // The field type in go, if any
+	GoReceiverName   string          // The name of method & var receiver in go, if any
+	GoFieldName      string          // The name of the method or var in go, if any
+	Args             []FieldArgument // A list of arguments to be passed to this field
+	ForceResolver    bool            // Should be emit Resolver method
+	MethodHasContext bool            // If this is bound to a go method, does the method also take a context
+	NoErr            bool            // If this is bound to a go method, does that method have an error as the second argument
+	Object           *Object         // A link back to the parent object
+	Default          interface{}     // The default value
 }
 
 type FieldArgument struct {
@@ -209,6 +210,10 @@ func (f *Field) CallArgs() string {
 		if !f.Object.Root {
 			args = append(args, "obj")
 		}
+	} else {
+		if f.MethodHasContext {
+			args = append(args, "ctx")
+		}
 	}
 
 	for _, arg := range f.Args {
diff --git a/codegen/util.go b/codegen/util.go
index 2ecd7c9e750..c6672bfc194 100644
--- a/codegen/util.go
+++ b/codegen/util.go
@@ -255,7 +255,19 @@ func bindMethod(imports *Imports, t types.Type, field *Field) error {
 	} else if sig.Results().Len() != 2 {
 		return fmt.Errorf("method has wrong number of args")
 	}
-	newArgs, err := matchArgs(field, sig.Params())
+	params := sig.Params()
+	// If the first argument is the context, remove it from the comparison and set
+	// the MethodHasContext flag so that the context will be passed to this model's method
+	if params.Len() > 0 && params.At(0).Type().String() == "context.Context" {
+		field.MethodHasContext = true
+		vars := make([]*types.Var, params.Len()-1)
+		for i := 1; i < params.Len(); i++ {
+			vars[i-1] = params.At(i)
+		}
+		params = types.NewTuple(vars...)
+	}
+
+	newArgs, err := matchArgs(field, params)
 	if err != nil {
 		return err
 	}