22
33import com .flit .protoc .gen .server .BaseGenerator ;
44import com .flit .protoc .gen .server .TypeMapper ;
5+ import com .flit .protoc .gen .server .Types ;
56import com .google .common .net .MediaType ;
67import com .google .protobuf .DescriptorProtos ;
78import com .google .protobuf .DescriptorProtos .MethodDescriptorProto ;
1011import com .squareup .javapoet .ClassName ;
1112import com .squareup .javapoet .FieldSpec ;
1213import com .squareup .javapoet .MethodSpec ;
14+ import com .squareup .javapoet .ParameterSpec ;
1315import com .squareup .javapoet .TypeSpec ;
1416import com .squareup .javapoet .TypeSpec .Builder ;
1517import java .util .Collections ;
@@ -22,6 +24,9 @@ public class RpcGenerator extends BaseGenerator {
2224 public static final ClassName POST = ClassName .bestGuess ("javax.ws.rs.POST" );
2325 public static final ClassName PRODUCES = ClassName .bestGuess ("javax.ws.rs.Produces" );
2426 public static final ClassName CONSUMES = ClassName .bestGuess ("javax.ws.rs.Consumes" );
27+ public static final ClassName CONTEXT = ClassName .bestGuess ("javax.ws.rs.core.Context" );
28+ public static final ClassName HttpServletRequest = ClassName .bestGuess ("javax.servlet.http.HttpServletRequest" );
29+ public static final ClassName HttpServletResponse = ClassName .bestGuess ("javax.servlet.http.HttpServletResponse" );
2530 private final String context ;
2631 private final Builder rpcResource ;
2732
@@ -49,23 +54,49 @@ private void addConstructor() {
4954
5055 private void addHandleMethod (MethodDescriptorProto mdp ) {
5156 ClassName inputType = mapper .get (mdp .getInputType ());
57+ ClassName outputType = mapper .get (mdp .getOutputType ());
5258 rpcResource .addMethod (MethodSpec .methodBuilder ("handle" + mdp .getName ())
5359 .addModifiers (Modifier .PUBLIC )
5460 .addAnnotation (POST )
5561 .addAnnotation (AnnotationSpec .builder (PATH )
5662 .addMember ("value" , "$S" , "/" + mdp .getName ())
5763 .build ())
58- .addAnnotation (AnnotationSpec .builder (PRODUCES )
59- .addMember ("value" , "$S" , MediaType .PROTOBUF .toString ())
60- .addMember ("value" , "$S" , MediaType .JSON_UTF_8 .toString ())
61- .build ())
62- .addAnnotation (AnnotationSpec .builder (CONSUMES )
63- .addMember ("value" , "$S" , MediaType .PROTOBUF .toString ())
64- .addMember ("value" , "$S" , MediaType .JSON_UTF_8 .toString ())
65- .build ())
66- .addParameter (inputType , "request" )
67- .addStatement ("return Response.ok(service.handle$L(request)).build()" , mdp .getName ())
68- .returns (ClassName .bestGuess ("javax.ws.rs.core.Response" ))
64+ .addParameter (ParameterSpec .builder (HttpServletRequest , "request" )
65+ .addAnnotation (CONTEXT ).build ())
66+ .addParameter (ParameterSpec .builder (HttpServletResponse , "response" )
67+ .addAnnotation (CONTEXT ).build ())
68+ .addException (Types .Exception )
69+ .addStatement ("boolean json = false" )
70+ .addStatement ("final $T data" , inputType )
71+ .beginControlFlow ("if (request.getContentType().equals($S))" , MediaType .PROTOBUF .toString ())
72+ .addStatement ("data = $T.parseFrom(request.getInputStream())" , inputType )
73+ .nextControlFlow ("else if (request.getContentType().startsWith($S))" , "application/json" )
74+ .addStatement ("json = true" )
75+ .addStatement ("$T.Builder builder = $T.newBuilder()" , inputType , inputType )
76+ .addStatement ("$T.parser().merge(new $T(request.getInputStream(), $T.UTF_8), builder)" ,
77+ Types .JsonFormat ,
78+ Types .InputStreamReader ,
79+ Types .StandardCharsets )
80+ .addStatement ("data = builder.build()" )
81+ .nextControlFlow ("else" )
82+ .addStatement ("response.setStatus(415)" )
83+ .addStatement ("response.flushBuffer()" )
84+ .addStatement ("return" )
85+ .endControlFlow ()
86+ // route to the service
87+ .addStatement ("$T retval = service.handle$L(data)" , outputType , mdp .getName ())
88+ .addStatement ("response.setStatus(200)" )
89+ // send the response
90+ .beginControlFlow ("if (json)" )
91+ .addStatement ("response.setContentType($S)" , MediaType .JSON_UTF_8 .toString ())
92+ .addStatement ("response.getOutputStream().write($T.printer().omittingInsignificantWhitespace().print(retval).getBytes($T.UTF_8))" ,
93+ Types .JsonFormat ,
94+ Types .StandardCharsets )
95+ .nextControlFlow ("else" )
96+ .addStatement ("response.setContentType($S)" , MediaType .PROTOBUF .toString ())
97+ .addStatement ("retval.writeTo(response.getOutputStream())" )
98+ .endControlFlow ()
99+ .addStatement ("response.flushBuffer()" )
69100 .build ());
70101 }
71102
0 commit comments