blob: 4400a1af7494b1d268b29808e19509f9853b229f [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 Haraszti2bdb6b32016-11-03 16:56:17 -070017
Zsolt Haraszti034db372016-10-03 22:26:41 -070018import sys
19
20from google.protobuf.compiler import plugin_pb2 as plugin
Zsolt Harasztibae12752016-10-10 09:55:30 -070021from google.protobuf.descriptor_pb2 import ServiceDescriptorProto, \
22 MethodOptions
Zsolt Haraszti034db372016-10-03 22:26:41 -070023from jinja2 import Template
24from simplejson import dumps
25
Zsolt Harasztibae12752016-10-10 09:55:30 -070026from chameleon.protos.third_party.google.api import annotations_pb2, http_pb2
27_ = annotations_pb2, http_pb2 # to keep import line from being optimized out
Zsolt Haraszti034db372016-10-03 22:26:41 -070028
29
30template = Template("""
31# Generated file; please do not edit
32
33from simplejson import dumps, load
34from structlog import get_logger
Zsolt Haraszti66862032016-11-28 14:28:39 -080035from google.protobuf.json_format import MessageToDict, ParseDict
Zsolt Haraszti00d9a842016-11-23 11:18:23 -080036from twisted.internet.defer import inlineCallbacks, returnValue
Zsolt Haraszti034db372016-10-03 22:26:41 -070037
38{% set package = file_name.replace('.proto', '') %}
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -080039
40{% for pypackage, module in includes %}
41{% if pypackage %}
42from {{ pypackage }} import {{ module }}
43{% else %}
44import {{ module }}
45{% endif %}
46{% endfor %}
Zsolt Haraszti034db372016-10-03 22:26:41 -070047
48log = get_logger()
49
50def add_routes(app, grpc_client):
51
Zsolt Haraszti023ea7c2016-10-16 19:30:34 -070052 pass # so that if no endpoints are defined, Python is still happy
53
Zsolt Haraszti034db372016-10-03 22:26:41 -070054 {% for method in methods %}
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -080055 {% set method_name = method['service'].rpartition('.')[2] + '_' + method['method'] %}
Zsolt Haraszti034db372016-10-03 22:26:41 -070056 {% set path = method['path'].replace('{', '<string:').replace('}', '>') %}
57 @app.route('{{ path }}', methods=['{{ method['verb'].upper() }}'])
Zsolt Haraszti00d9a842016-11-23 11:18:23 -080058 @inlineCallbacks
Zsolt Haraszti034db372016-10-03 22:26:41 -070059 def {{ method_name }}(server, request, **kw):
60 log.debug('{{ method_name }}', request=request, server=server, **kw)
61 {% if method['body'] == '*' %}
62 data = load(request.content)
63 data.update(kw)
64 {% elif method['body'] == '' %}
65 data = kw
66 {% else %}
Zsolt Haraszti66862032016-11-28 14:28:39 -080067 raise NotImplementedError('cannot handle specific body field list')
Zsolt Haraszti034db372016-10-03 22:26:41 -070068 {% endif %}
Zsolt Haraszti2bdb6b32016-11-03 16:56:17 -070069 try:
Zsolt Haraszti66862032016-11-28 14:28:39 -080070 req = ParseDict(data, {{ type_map[method['input_type']] }}())
Zsolt Haraszti2bdb6b32016-11-03 16:56:17 -070071 except Exception, e:
72 log.error('cannot-convert-to-protobuf', e=e, data=data)
73 raise
Zsolt Haraszti66862032016-11-28 14:28:39 -080074 res, metadata = yield grpc_client.invoke(
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -080075 {{ type_map[method['service']] }}Stub,
Zsolt Haraszti66862032016-11-28 14:28:39 -080076 '{{ method['method'] }}', req, request.getAllHeaders().items())
Zsolt Haraszti2bdb6b32016-11-03 16:56:17 -070077 try:
Zsolt Haraszti00d9a842016-11-23 11:18:23 -080078 out_data = MessageToDict(res, True, True)
Zsolt Haraszti2bdb6b32016-11-03 16:56:17 -070079 except AttributeError, e:
80 filename = '/tmp/chameleon_failed_to_convert_data.pbd'
81 with file(filename, 'w') as f:
82 f.write(res.SerializeToString())
83 log.error('cannot-convert-from-protobuf', outdata_saved=filename)
84 raise
Zsolt Haraszti66862032016-11-28 14:28:39 -080085 for key, value in metadata:
86 request.setHeader(key, value)
Zsolt Haraszti034db372016-10-03 22:26:41 -070087 request.setHeader('Content-Type', 'application/json')
88 log.debug('{{ method_name }}', **out_data)
Zsolt Haraszti00d9a842016-11-23 11:18:23 -080089 returnValue(dumps(out_data))
Zsolt Haraszti034db372016-10-03 22:26:41 -070090
91 {% endfor %}
92
93""", trim_blocks=True, lstrip_blocks=True)
94
95
96def traverse_methods(proto_file):
97
98 package = proto_file.name
99 for service in proto_file.service:
100 assert isinstance(service, ServiceDescriptorProto)
101
102 for method in service.method:
103 options = method.options
104 assert isinstance(options, MethodOptions)
105 for fd, http in options.ListFields():
106 if fd.full_name == 'google.api.http':
107 assert fd.name == 'http'
108 assert isinstance(http, http_pb2.HttpRule)
109
110 input_type = method.input_type
111 if input_type.startswith('.'):
112 input_type = input_type[1:]
113
114 output_type = method.output_type
115 if output_type.startswith('.'):
116 output_type = output_type[1:]
117
118 if http.delete:
119 verb = 'delete'
120 path = http.delete
121 elif http.get:
122 verb = 'get'
123 path = http.get
124 elif http.patch:
125 verb = 'patch'
126 path = http.patch
127 elif http.post:
128 verb = 'post'
129 path = http.post
130 elif http.put:
131 verb = 'put'
132 path = http.put
133 else:
134 raise AttributeError('No valid verb in method %s' %
135 method.name)
136
137 body = http.body
138
139 data = {
140 'package': package,
141 'filename': proto_file.name,
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800142 'service': proto_file.package + '.' + service.name,
Zsolt Haraszti034db372016-10-03 22:26:41 -0700143 'method': method.name,
144 'input_type': input_type,
145 'output_type': output_type,
146 'path': path,
147 'verb': verb,
148 'body': body
149 }
150
151 yield data
152
153
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800154def generate_gw_code(file_name, methods, type_map, includes):
155 return template.render(file_name=file_name, methods=methods,
156 type_map=type_map, includes=includes)
157
158
159class IncludeManager(object):
160 # need to keep track of what files define what message types and
161 # under what package name. Later, when we analyze the methods, we
162 # need to be able to derive the list of files we need to load and we
163 # also need to replce the <proto-package-name>.<artifact-name> in the
164 # templates with <python-package-name>.<artifact-name> so Python can
165 # resolve these.
166 def __init__(self):
167 self.package_to_localname = {}
168 self.fullname_to_filename = {}
169 self.prefix_table = [] # sorted table of top-level symbols in protos
170 self.type_map = {} # full name as used in .proto -> python name
171 self.includes_needed = set() # names of files needed to be included
172 self.filename_to_module = {} # filename -> (package, module)
173
174 def extend_symbol_tables(self, proto_file):
175 # keep track of what file adds what top-level symbol to what abstract
176 # package name
177 package_name = proto_file.package
178 file_name = proto_file.name
179 self._add_filename(file_name)
180 all_defs = list(proto_file.message_type)
181 all_defs.extend(list(proto_file.enum_type))
182 all_defs.extend(list(proto_file.service))
183 for typedef in all_defs:
184 name = typedef.name
185 fullname = package_name + '.' + name
186 self.fullname_to_filename[fullname] = file_name
187 self.package_to_localname.setdefault(package_name, []).append(name)
188 self._update_prefix_table()
189
190 def _add_filename(self, filename):
191 if filename not in self.filename_to_module:
192 python_path = filename.replace('.proto', '_pb2').replace('/', '.')
193 package_name, _, module_name = python_path.rpartition('.')
194 self.filename_to_module[filename] = (package_name, module_name)
195
196 def _update_prefix_table(self):
197 # make a sorted list symbol prefixes needed to resolv for potential use
198 # of nested symbols
199 self.prefix_table = sorted(self.fullname_to_filename.iterkeys(),
200 reverse=True)
201
202 def _find_matching_prefix(self, fullname):
203 for prefix in self.prefix_table:
204 if fullname.startswith(prefix):
205 return prefix
206 # This should never happen
207 raise Exception('No match for type name "{}"'.format(fullname))
208
209 def add_needed_symbol(self, fullname):
210 if fullname in self.type_map:
211 return
212 top_level_symbol = self._find_matching_prefix(fullname)
213 name = top_level_symbol.rpartition('.')[2]
214 nested_name = fullname[len(top_level_symbol):] # may be empty
215 file_name = self.fullname_to_filename[top_level_symbol]
216 self.includes_needed.add(file_name)
217 module_name = self.filename_to_module[file_name][1]
218 python_name = module_name + '.' + name + nested_name
219 self.type_map[fullname] = python_name
220
221 def get_type_map(self):
222 return self.type_map
223
224 def get_includes(self):
225 return sorted(
226 self.filename_to_module[fn] for fn in self.includes_needed)
Zsolt Haraszti034db372016-10-03 22:26:41 -0700227
228
229def generate_code(request, response):
230
231 assert isinstance(request, plugin.CodeGeneratorRequest)
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800232
233 include_manager = IncludeManager()
Zsolt Haraszti034db372016-10-03 22:26:41 -0700234 for proto_file in request.proto_file:
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800235
236 include_manager.extend_symbol_tables(proto_file)
237
238 methods = []
Zsolt Haraszti034db372016-10-03 22:26:41 -0700239
240 for data in traverse_methods(proto_file):
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800241 methods.append(data)
242 include_manager.add_needed_symbol(data['input_type'])
243 include_manager.add_needed_symbol(data['output_type'])
244 include_manager.add_needed_symbol(data['service'])
245
246 type_map = include_manager.get_type_map()
247 includes = include_manager.get_includes()
Zsolt Haraszti034db372016-10-03 22:26:41 -0700248
249 # as a nice side-effect, generate a json file capturing the essence
250 # of the RPC method entries
251 f = response.file.add()
252 f.name = proto_file.name + '.json'
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800253 f.content = dumps(dict(
254 type_rename_map=type_map,
255 includes=includes,
256 methods=methods), indent=4)
Zsolt Haraszti034db372016-10-03 22:26:41 -0700257
258 # generate the real Python code file
259 f = response.file.add()
260 assert proto_file.name.endswith('.proto')
261 f.name = proto_file.name.replace('.proto', '_gw.py')
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800262 f.content = generate_gw_code(proto_file.name,
263 methods, type_map, includes)
Zsolt Haraszti034db372016-10-03 22:26:41 -0700264
265
266if __name__ == '__main__':
267
268 if len(sys.argv) >= 2:
269 # read input from file, to allow troubleshooting
270 with open(sys.argv[1], 'r') as f:
271 data = f.read()
272 else:
273 # read input from stdin
274 data = sys.stdin.read()
Zsolt Haraszti7eeb2b32016-11-06 14:04:55 -0800275 # with file('/tmp/buf', 'wb') as f:
276 # f.write(data)
Zsolt Haraszti034db372016-10-03 22:26:41 -0700277
278 # parse request
279 request = plugin.CodeGeneratorRequest()
280 request.ParseFromString(data)
281
282 # create response object
283 response = plugin.CodeGeneratorResponse()
284
285 # generate the output and the response
286 generate_code(request, response)
287
288 # serialize the response
289 output = response.SerializeToString()
290
291 # write response to stdout
292 sys.stdout.write(output)