Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 1 | #!/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 Haraszti | 2bdb6b3 | 2016-11-03 16:56:17 -0700 | [diff] [blame] | 17 | |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 18 | import sys |
| 19 | |
| 20 | from google.protobuf.compiler import plugin_pb2 as plugin |
Zsolt Haraszti | bae1275 | 2016-10-10 09:55:30 -0700 | [diff] [blame] | 21 | from google.protobuf.descriptor_pb2 import ServiceDescriptorProto, \ |
| 22 | MethodOptions |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 23 | from jinja2 import Template |
| 24 | from simplejson import dumps |
| 25 | |
Zsolt Haraszti | bae1275 | 2016-10-10 09:55:30 -0700 | [diff] [blame] | 26 | from 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 Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 28 | |
| 29 | |
| 30 | template = Template(""" |
| 31 | # Generated file; please do not edit |
| 32 | |
| 33 | from simplejson import dumps, load |
| 34 | from structlog import get_logger |
| 35 | from protobuf_to_dict import protobuf_to_dict, dict_to_protobuf |
| 36 | |
| 37 | {% set package = file_name.replace('.proto', '') %} |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 38 | |
| 39 | {% for pypackage, module in includes %} |
| 40 | {% if pypackage %} |
| 41 | from {{ pypackage }} import {{ module }} |
| 42 | {% else %} |
| 43 | import {{ module }} |
| 44 | {% endif %} |
| 45 | {% endfor %} |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 46 | |
| 47 | log = get_logger() |
| 48 | |
| 49 | def add_routes(app, grpc_client): |
| 50 | |
Zsolt Haraszti | 023ea7c | 2016-10-16 19:30:34 -0700 | [diff] [blame] | 51 | pass # so that if no endpoints are defined, Python is still happy |
| 52 | |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 53 | {% for method in methods %} |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 54 | {% set method_name = method['service'].rpartition('.')[2] + '_' + method['method'] %} |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 55 | {% set path = method['path'].replace('{', '<string:').replace('}', '>') %} |
| 56 | @app.route('{{ path }}', methods=['{{ method['verb'].upper() }}']) |
| 57 | def {{ method_name }}(server, request, **kw): |
| 58 | log.debug('{{ method_name }}', request=request, server=server, **kw) |
| 59 | {% if method['body'] == '*' %} |
| 60 | data = load(request.content) |
| 61 | data.update(kw) |
| 62 | {% elif method['body'] == '' %} |
| 63 | data = kw |
| 64 | {% else %} |
| 65 | riase NotImplementedError('cannot handle specific body field list') |
| 66 | {% endif %} |
Zsolt Haraszti | 2bdb6b3 | 2016-11-03 16:56:17 -0700 | [diff] [blame] | 67 | try: |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 68 | req = dict_to_protobuf({{ type_map[method['input_type']] }}, data) |
Zsolt Haraszti | 2bdb6b3 | 2016-11-03 16:56:17 -0700 | [diff] [blame] | 69 | except Exception, e: |
| 70 | log.error('cannot-convert-to-protobuf', e=e, data=data) |
| 71 | raise |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 72 | res = grpc_client.invoke( |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 73 | {{ type_map[method['service']] }}Stub, |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 74 | '{{ method['method'] }}', req) |
Zsolt Haraszti | 2bdb6b3 | 2016-11-03 16:56:17 -0700 | [diff] [blame] | 75 | try: |
| 76 | out_data = protobuf_to_dict(res, use_enum_labels=True) |
| 77 | except AttributeError, e: |
| 78 | filename = '/tmp/chameleon_failed_to_convert_data.pbd' |
| 79 | with file(filename, 'w') as f: |
| 80 | f.write(res.SerializeToString()) |
| 81 | log.error('cannot-convert-from-protobuf', outdata_saved=filename) |
| 82 | raise |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 83 | request.setHeader('Content-Type', 'application/json') |
| 84 | log.debug('{{ method_name }}', **out_data) |
| 85 | return dumps(out_data) |
| 86 | |
| 87 | {% endfor %} |
| 88 | |
| 89 | """, trim_blocks=True, lstrip_blocks=True) |
| 90 | |
| 91 | |
| 92 | def traverse_methods(proto_file): |
| 93 | |
| 94 | package = proto_file.name |
| 95 | for service in proto_file.service: |
| 96 | assert isinstance(service, ServiceDescriptorProto) |
| 97 | |
| 98 | for method in service.method: |
| 99 | options = method.options |
| 100 | assert isinstance(options, MethodOptions) |
| 101 | for fd, http in options.ListFields(): |
| 102 | if fd.full_name == 'google.api.http': |
| 103 | assert fd.name == 'http' |
| 104 | assert isinstance(http, http_pb2.HttpRule) |
| 105 | |
| 106 | input_type = method.input_type |
| 107 | if input_type.startswith('.'): |
| 108 | input_type = input_type[1:] |
| 109 | |
| 110 | output_type = method.output_type |
| 111 | if output_type.startswith('.'): |
| 112 | output_type = output_type[1:] |
| 113 | |
| 114 | if http.delete: |
| 115 | verb = 'delete' |
| 116 | path = http.delete |
| 117 | elif http.get: |
| 118 | verb = 'get' |
| 119 | path = http.get |
| 120 | elif http.patch: |
| 121 | verb = 'patch' |
| 122 | path = http.patch |
| 123 | elif http.post: |
| 124 | verb = 'post' |
| 125 | path = http.post |
| 126 | elif http.put: |
| 127 | verb = 'put' |
| 128 | path = http.put |
| 129 | else: |
| 130 | raise AttributeError('No valid verb in method %s' % |
| 131 | method.name) |
| 132 | |
| 133 | body = http.body |
| 134 | |
| 135 | data = { |
| 136 | 'package': package, |
| 137 | 'filename': proto_file.name, |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 138 | 'service': proto_file.package + '.' + service.name, |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 139 | 'method': method.name, |
| 140 | 'input_type': input_type, |
| 141 | 'output_type': output_type, |
| 142 | 'path': path, |
| 143 | 'verb': verb, |
| 144 | 'body': body |
| 145 | } |
| 146 | |
| 147 | yield data |
| 148 | |
| 149 | |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 150 | def generate_gw_code(file_name, methods, type_map, includes): |
| 151 | return template.render(file_name=file_name, methods=methods, |
| 152 | type_map=type_map, includes=includes) |
| 153 | |
| 154 | |
| 155 | class IncludeManager(object): |
| 156 | # need to keep track of what files define what message types and |
| 157 | # under what package name. Later, when we analyze the methods, we |
| 158 | # need to be able to derive the list of files we need to load and we |
| 159 | # also need to replce the <proto-package-name>.<artifact-name> in the |
| 160 | # templates with <python-package-name>.<artifact-name> so Python can |
| 161 | # resolve these. |
| 162 | def __init__(self): |
| 163 | self.package_to_localname = {} |
| 164 | self.fullname_to_filename = {} |
| 165 | self.prefix_table = [] # sorted table of top-level symbols in protos |
| 166 | self.type_map = {} # full name as used in .proto -> python name |
| 167 | self.includes_needed = set() # names of files needed to be included |
| 168 | self.filename_to_module = {} # filename -> (package, module) |
| 169 | |
| 170 | def extend_symbol_tables(self, proto_file): |
| 171 | # keep track of what file adds what top-level symbol to what abstract |
| 172 | # package name |
| 173 | package_name = proto_file.package |
| 174 | file_name = proto_file.name |
| 175 | self._add_filename(file_name) |
| 176 | all_defs = list(proto_file.message_type) |
| 177 | all_defs.extend(list(proto_file.enum_type)) |
| 178 | all_defs.extend(list(proto_file.service)) |
| 179 | for typedef in all_defs: |
| 180 | name = typedef.name |
| 181 | fullname = package_name + '.' + name |
| 182 | self.fullname_to_filename[fullname] = file_name |
| 183 | self.package_to_localname.setdefault(package_name, []).append(name) |
| 184 | self._update_prefix_table() |
| 185 | |
| 186 | def _add_filename(self, filename): |
| 187 | if filename not in self.filename_to_module: |
| 188 | python_path = filename.replace('.proto', '_pb2').replace('/', '.') |
| 189 | package_name, _, module_name = python_path.rpartition('.') |
| 190 | self.filename_to_module[filename] = (package_name, module_name) |
| 191 | |
| 192 | def _update_prefix_table(self): |
| 193 | # make a sorted list symbol prefixes needed to resolv for potential use |
| 194 | # of nested symbols |
| 195 | self.prefix_table = sorted(self.fullname_to_filename.iterkeys(), |
| 196 | reverse=True) |
| 197 | |
| 198 | def _find_matching_prefix(self, fullname): |
| 199 | for prefix in self.prefix_table: |
| 200 | if fullname.startswith(prefix): |
| 201 | return prefix |
| 202 | # This should never happen |
| 203 | raise Exception('No match for type name "{}"'.format(fullname)) |
| 204 | |
| 205 | def add_needed_symbol(self, fullname): |
| 206 | if fullname in self.type_map: |
| 207 | return |
| 208 | top_level_symbol = self._find_matching_prefix(fullname) |
| 209 | name = top_level_symbol.rpartition('.')[2] |
| 210 | nested_name = fullname[len(top_level_symbol):] # may be empty |
| 211 | file_name = self.fullname_to_filename[top_level_symbol] |
| 212 | self.includes_needed.add(file_name) |
| 213 | module_name = self.filename_to_module[file_name][1] |
| 214 | python_name = module_name + '.' + name + nested_name |
| 215 | self.type_map[fullname] = python_name |
| 216 | |
| 217 | def get_type_map(self): |
| 218 | return self.type_map |
| 219 | |
| 220 | def get_includes(self): |
| 221 | return sorted( |
| 222 | self.filename_to_module[fn] for fn in self.includes_needed) |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 223 | |
| 224 | |
| 225 | def generate_code(request, response): |
| 226 | |
| 227 | assert isinstance(request, plugin.CodeGeneratorRequest) |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 228 | |
| 229 | include_manager = IncludeManager() |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 230 | for proto_file in request.proto_file: |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 231 | |
| 232 | include_manager.extend_symbol_tables(proto_file) |
| 233 | |
| 234 | methods = [] |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 235 | |
| 236 | for data in traverse_methods(proto_file): |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 237 | methods.append(data) |
| 238 | include_manager.add_needed_symbol(data['input_type']) |
| 239 | include_manager.add_needed_symbol(data['output_type']) |
| 240 | include_manager.add_needed_symbol(data['service']) |
| 241 | |
| 242 | type_map = include_manager.get_type_map() |
| 243 | includes = include_manager.get_includes() |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 244 | |
| 245 | # as a nice side-effect, generate a json file capturing the essence |
| 246 | # of the RPC method entries |
| 247 | f = response.file.add() |
| 248 | f.name = proto_file.name + '.json' |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 249 | f.content = dumps(dict( |
| 250 | type_rename_map=type_map, |
| 251 | includes=includes, |
| 252 | methods=methods), indent=4) |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 253 | |
| 254 | # generate the real Python code file |
| 255 | f = response.file.add() |
| 256 | assert proto_file.name.endswith('.proto') |
| 257 | f.name = proto_file.name.replace('.proto', '_gw.py') |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 258 | f.content = generate_gw_code(proto_file.name, |
| 259 | methods, type_map, includes) |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 260 | |
| 261 | |
| 262 | if __name__ == '__main__': |
| 263 | |
| 264 | if len(sys.argv) >= 2: |
| 265 | # read input from file, to allow troubleshooting |
| 266 | with open(sys.argv[1], 'r') as f: |
| 267 | data = f.read() |
| 268 | else: |
| 269 | # read input from stdin |
| 270 | data = sys.stdin.read() |
Zsolt Haraszti | 7eeb2b3 | 2016-11-06 14:04:55 -0800 | [diff] [blame] | 271 | # with file('/tmp/buf', 'wb') as f: |
| 272 | # f.write(data) |
Zsolt Haraszti | 034db37 | 2016-10-03 22:26:41 -0700 | [diff] [blame] | 273 | |
| 274 | # parse request |
| 275 | request = plugin.CodeGeneratorRequest() |
| 276 | request.ParseFromString(data) |
| 277 | |
| 278 | # create response object |
| 279 | response = plugin.CodeGeneratorResponse() |
| 280 | |
| 281 | # generate the output and the response |
| 282 | generate_code(request, response) |
| 283 | |
| 284 | # serialize the response |
| 285 | output = response.SerializeToString() |
| 286 | |
| 287 | # write response to stdout |
| 288 | sys.stdout.write(output) |