blob: 5398c32eef51717fc3ca92467efd627f5bd97217 [file] [log] [blame]
Zsolt Haraszti034db372016-10-03 22:26:41 -07001#!/usr/bin/env python
2#
3# Copyright 2016 the original author or authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
Zsolt Haraszti034db372016-10-03 22:26:41 -070017import sys
18
19from google.protobuf.compiler import plugin_pb2 as plugin
Zsolt Harasztibae12752016-10-10 09:55:30 -070020from google.protobuf.descriptor_pb2 import ServiceDescriptorProto, \
21 MethodOptions
Zsolt Haraszti034db372016-10-03 22:26:41 -070022from jinja2 import Template
23from simplejson import dumps
24
Zsolt Harasztibae12752016-10-10 09:55:30 -070025from chameleon.protos.third_party.google.api import annotations_pb2, http_pb2
26_ = annotations_pb2, http_pb2 # to keep import line from being optimized out
Zsolt Haraszti034db372016-10-03 22:26:41 -070027
28
29template = Template("""
30# Generated file; please do not edit
31
32from simplejson import dumps, load
33from structlog import get_logger
34from protobuf_to_dict import protobuf_to_dict, dict_to_protobuf
35
36{% set package = file_name.replace('.proto', '') %}
37import {{ package + '_pb2' }} as {{ package }}
38
39log = get_logger()
40
41def add_routes(app, grpc_client):
42
Zsolt Haraszti023ea7c2016-10-16 19:30:34 -070043 pass # so that if no endpoints are defined, Python is still happy
44
Zsolt Haraszti034db372016-10-03 22:26:41 -070045 {% for method in methods %}
46 {% set method_name = method['service'] + '_' + method['method'] %}
47 {% set path = method['path'].replace('{', '<string:').replace('}', '>') %}
48 @app.route('{{ path }}', methods=['{{ method['verb'].upper() }}'])
49 def {{ method_name }}(server, request, **kw):
50 log.debug('{{ method_name }}', request=request, server=server, **kw)
51 {% if method['body'] == '*' %}
52 data = load(request.content)
53 data.update(kw)
54 {% elif method['body'] == '' %}
55 data = kw
56 {% else %}
57 riase NotImplementedError('cannot handle specific body field list')
58 {% endif %}
59 req = dict_to_protobuf({{ method['input_type'] }}, data)
60 res = grpc_client.invoke(
61 {{ '.'.join([package, method['service']]) }}Stub,
62 '{{ method['method'] }}', req)
63 out_data = protobuf_to_dict(res, use_enum_labels=True)
64 request.setHeader('Content-Type', 'application/json')
65 log.debug('{{ method_name }}', **out_data)
66 return dumps(out_data)
67
68 {% endfor %}
69
70""", trim_blocks=True, lstrip_blocks=True)
71
72
73def traverse_methods(proto_file):
74
75 package = proto_file.name
76 for service in proto_file.service:
77 assert isinstance(service, ServiceDescriptorProto)
78
79 for method in service.method:
80 options = method.options
81 assert isinstance(options, MethodOptions)
82 for fd, http in options.ListFields():
83 if fd.full_name == 'google.api.http':
84 assert fd.name == 'http'
85 assert isinstance(http, http_pb2.HttpRule)
86
87 input_type = method.input_type
88 if input_type.startswith('.'):
89 input_type = input_type[1:]
90
91 output_type = method.output_type
92 if output_type.startswith('.'):
93 output_type = output_type[1:]
94
95 if http.delete:
96 verb = 'delete'
97 path = http.delete
98 elif http.get:
99 verb = 'get'
100 path = http.get
101 elif http.patch:
102 verb = 'patch'
103 path = http.patch
104 elif http.post:
105 verb = 'post'
106 path = http.post
107 elif http.put:
108 verb = 'put'
109 path = http.put
110 else:
111 raise AttributeError('No valid verb in method %s' %
112 method.name)
113
114 body = http.body
115
116 data = {
117 'package': package,
118 'filename': proto_file.name,
119 'service': service.name,
120 'method': method.name,
121 'input_type': input_type,
122 'output_type': output_type,
123 'path': path,
124 'verb': verb,
125 'body': body
126 }
127
128 yield data
129
130
131def generate_gw_code(file_name, methods):
132 return template.render(file_name=file_name, methods=methods)
133
134
135def generate_code(request, response):
136
137 assert isinstance(request, plugin.CodeGeneratorRequest)
138 for proto_file in request.proto_file:
139 output = []
140
141 for data in traverse_methods(proto_file):
142 output.append(data)
143
144 # as a nice side-effect, generate a json file capturing the essence
145 # of the RPC method entries
146 f = response.file.add()
147 f.name = proto_file.name + '.json'
148 f.content = dumps(output, indent=4)
149
150 # generate the real Python code file
151 f = response.file.add()
152 assert proto_file.name.endswith('.proto')
153 f.name = proto_file.name.replace('.proto', '_gw.py')
154 f.content = generate_gw_code(proto_file.name, output)
155
156
157if __name__ == '__main__':
158
159 if len(sys.argv) >= 2:
160 # read input from file, to allow troubleshooting
161 with open(sys.argv[1], 'r') as f:
162 data = f.read()
163 else:
164 # read input from stdin
165 data = sys.stdin.read()
166
167 # parse request
168 request = plugin.CodeGeneratorRequest()
169 request.ParseFromString(data)
170
171 # create response object
172 response = plugin.CodeGeneratorResponse()
173
174 # generate the output and the response
175 generate_code(request, response)
176
177 # serialize the response
178 output = response.SerializeToString()
179
180 # write response to stdout
181 sys.stdout.write(output)