Propagates context cancellation.

mapper/mapper_test.go

@@ -278,6 +278,24 @@ func TestFailedStreamingCallbacks(t *testing.T) {
 	protoResult("testMappingClient.BlogsCF", sResp, err, sErr, true)
 }
 
+func TestCanceledContext(t *testing.T) {
+	// todo defer
+	timeout, _ := time.ParseDuration("5s")
+
+	// unary
+	ctx, cancelFunc := context.WithCancel(context.TODO())
+	time.AfterFunc(timeout, cancelFunc)
+	resp, sErr := testMappingClient.CanceledUnaryContext(ctx, &td.EmptyRequest{})
+	protoResult("testMappingClient.CanceledUnaryContext", resp, nil, sErr, true) // canceled during query execution
+
+	// streaming
+	ctx, cancelFunc = context.WithCancel(context.TODO())
+	time.AfterFunc(timeout, cancelFunc)
+	stream, err := testMappingClient.CanceledStreamContext(ctx, &td.EmptyRequest{})
+	sResp, sErr := postReader(stream)
+	protoResult("testMappingClient.CanceledStreamContext", sResp, err, sErr, true)
+}
+
 func blogStreamReader(stream ex.BlogQueryService_SelectBlogsClient) ([]proto.Message, error) {
 	var responses []proto.Message
 	for {

plugin/imports.go

@@ -35,6 +35,7 @@ func (p *SqlPlugin) setStreamingImports() {
 	p.Pkg["codes"] = true
 	p.Pkg["status"] = true
 	p.Pkg["log"] = true
+	p.Pkg["context"] = true
 }
 
 func (p *SqlPlugin) setUnaryImports() {

templates/streaing_response.go

@@ -55,10 +55,12 @@ func (m *{{ .ServiceName }}MapServer) {{ .MethodName }}(r *{{ .RequestName }}, s
 	preparedSql, args, err := mapper.PrepareQuery(m.Dialect, sqlBuffer.Bytes())
 	if err != nil {
 		log.Printf("error preparing sql query.\n {{ .RequestName }} request: %s \n query: %s \n error: %s", r, rawSql, err)
-		return status.Error(codes.InvalidArgument, "request generated malformed query")
+		return status.Error(codes.InvalidArgument, "Request generated malformed query.")
 	}
-	rows, err := m.DB.Query(preparedSql, args...)
-	if err != nil {
+	rows, err := m.DB.QueryContext(stream.Context(), preparedSql, args...)
+	if stream.Context().Err() == context.Canceled {
+		return status.Error(codes.Canceled, "Client cancelled.")
+	} else if err != nil {
 		log.Printf("error executing query.\n {{ .RequestName }} request: %s \n query: %s \n error: %s", r, rawSql, err)
 		return status.Error(codes.Internal, err.Error())
 	} else {
@@ -69,8 +71,8 @@ func (m *{{ .ServiceName }}MapServer) {{ .MethodName }}(r *{{ .RequestName }}, s
 		m.{{ .MapperName }}Mapper, err =  mapper.New("{{ .MethodName }}", rows, &{{ .ResponseName }}{})
 		m.mapperGenMux.Unlock()
 		if err != nil {
-			log.Printf("error generating {{ .MapperName }}Mapper: %s", err)
-			return status.Error(codes.Internal, "error generating {{ .ResponseName }} mapping")
+			log.Printf("Error generating {{ .MapperName }}Mapper: %s", err)
+			return status.Error(codes.Internal, "Error generating {{ .ResponseName }} mapping.")
 		}
 		m.{{ .MapperName }}Mapper.Log()
 	}

templates/unary_response.go

@@ -52,8 +52,10 @@ func (m *{{ .ServiceName }}MapServer) {{ .MethodName }}(ctx context.Context, r *
 		return nil, status.Error(codes.InvalidArgument, "request generated malformed query")
 	}
 	{{- if eq .QueryType "Exec" }}
-	_, err = m.DB.Exec(preparedSql, args...)
-	if err != nil {
+	_, err = m.DB.ExecContext(ctx, preparedSql, args...)
+	if ctx.Err() == context.Canceled {
+		return nil, status.Error(codes.Canceled, "Client cancelled.")
+	} else if err != nil {
 		log.Printf("error executing query.\n {{ .RequestName }} request: %s \n query: %s \n error: %s", r, preparedSql, err)
 		return nil, status.Error(codes.InvalidArgument, "request generated malformed query")
 	}
@@ -66,8 +68,10 @@ func (m *{{ .ServiceName }}MapServer) {{ .MethodName }}(ctx context.Context, r *
 	resp :={{ .ResponseName }}{}
         return &resp, nil
 	{{ else if eq .QueryType "Query" }}
-	rows, err := m.DB.Query(preparedSql, args...)
-	if err != nil {
+	rows, err := m.DB.QueryContext(ctx, preparedSql, args...)
+	if ctx.Err() == context.Canceled {
+		return nil, status.Error(codes.Canceled, "Client cancelled.")
+	} else if err != nil {
 		log.Printf("error executing query.\n {{ .RequestName }} request: %s \n,query: %s \n error: %s", r, preparedSql, err)
 		return nil, status.Error(codes.InvalidArgument, "request generated malformed query")
 	} else {