entc/gen: merge context.go contents to ent.go file (ent#3355)

entc/gen/template.go

@@ -103,10 +103,6 @@ var (
 				"dialect/*/query/fields/init/*",
 			},
 		},
-		{
-			Name:   "context",
-			Format: "context.go",
-		},
 		{
 			Name:   "tx",
 			Format: "tx.go",
@@ -168,7 +164,7 @@ var (
 		},
 	}
 	// template files that were deleted and should be removed by the codegen.
-	deletedTemplates = []string{"config.go"}
+	deletedTemplates = []string{"config.go", "context.go"}
 	// patterns for extending partial-templates (included by other templates).
 	partialPatterns = [...]string{
 		"client/additional/*",

entc/gen/template/base.tmpl

@@ -39,6 +39,32 @@ type (
 	MutateFunc    = ent.MutateFunc
 )
 
+type clientCtxKey struct{}
+
+// FromContext returns a Client stored inside a context, or nil if there isn't one.
+func FromContext(ctx context.Context) *Client {
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
+	return c
+}
+
+// NewContext returns a new context with the given Client attached.
+func NewContext(parent context.Context, c *Client) context.Context {
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct {}
+
+// TxFromContext returns a Tx stored inside a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Tx attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
+}
+
 {{ $tmpl := printf "dialect/%s/order/signature" $.Storage }}
 {{ xtemplate $tmpl . }}