blob: 20249c2cfa7aae7df49cffd89e8ead520a8893bf [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
35from protobuf_to_dict import protobuf_to_dict, dict_to_protobuf
36
37{% set package = file_name.replace('.proto', '') %}
38import {{ package + '_pb2' }} as {{ package }}
39
40log = get_logger()
41
42def add_routes(app, grpc_client):
43
Zsolt Haraszti023ea7c2016-10-16 19:30:34 -070044 pass # so that if no endpoints are defined, Python is still happy
45
Zsolt Haraszti034db372016-10-03 22:26:41 -070046 {% for method in methods %}
47 {% set method_name = method['service'] + '_' + method['method'] %}
48 {% set path = method['path'].replace('{', '<string:').replace('}', '>') %}
49 @app.route('{{ path }}', methods=['{{ method['verb'].upper() }}'])
50 def {{ method_name }}(server, request, **kw):
51 log.debug('{{ method_name }}', request=request, server=server, **kw)
52 {% if method['body'] == '*' %}
53 data = load(request.content)
54 data.update(kw)
55 {% elif method['body'] == '' %}
56 data = kw
57 {% else %}
58 riase NotImplementedError('cannot handle specific body field list')
59 {% endif %}
Zsolt Haraszti2bdb6b32016-11-03 16:56:17 -070060 try:
61 req = dict_to_protobuf({{ method['input_type'] }}, data)
62 except Exception, e:
63 log.error('cannot-convert-to-protobuf', e=e, data=data)
64 raise
Zsolt Haraszti034db372016-10-03 22:26:41 -070065 res = grpc_client.invoke(
66 {{ '.'.join([package, method['service']]) }}Stub,
67 '{{ method['method'] }}', req)
Zsolt Haraszti2bdb6b32016-11-03 16:56:17 -070068 try:
69 out_data = protobuf_to_dict(res, use_enum_labels=True)
70 except AttributeError, e:
71 filename = '/tmp/chameleon_failed_to_convert_data.pbd'
72 with file(filename, 'w') as f:
73 f.write(res.SerializeToString())
74 log.error('cannot-convert-from-protobuf', outdata_saved=filename)
75 raise
Zsolt Haraszti034db372016-10-03 22:26:41 -070076 request.setHeader('Content-Type', 'application/json')
77 log.debug('{{ method_name }}', **out_data)
78 return dumps(out_data)
79
80 {% endfor %}
81
82""", trim_blocks=True, lstrip_blocks=True)
83
84
85def traverse_methods(proto_file):
86
87 package = proto_file.name
88 for service in proto_file.service:
89 assert isinstance(service, ServiceDescriptorProto)
90
91 for method in service.method:
92 options = method.options
93 assert isinstance(options, MethodOptions)
94 for fd, http in options.ListFields():
95 if fd.full_name == 'google.api.http':
96 assert fd.name == 'http'
97 assert isinstance(http, http_pb2.HttpRule)
98
99 input_type = method.input_type
100 if input_type.startswith('.'):
101 input_type = input_type[1:]
102
103 output_type = method.output_type
104 if output_type.startswith('.'):
105 output_type = output_type[1:]
106
107 if http.delete:
108 verb = 'delete'
109 path = http.delete
110 elif http.get:
111 verb = 'get'
112 path = http.get
113 elif http.patch:
114 verb = 'patch'
115 path = http.patch
116 elif http.post:
117 verb = 'post'
118 path = http.post
119 elif http.put:
120 verb = 'put'
121 path = http.put
122 else:
123 raise AttributeError('No valid verb in method %s' %
124 method.name)
125
126 body = http.body
127
128 data = {
129 'package': package,
130 'filename': proto_file.name,
131 'service': service.name,
132 'method': method.name,
133 'input_type': input_type,
134 'output_type': output_type,
135 'path': path,
136 'verb': verb,
137 'body': body
138 }
139
140 yield data
141
142
143def generate_gw_code(file_name, methods):
144 return template.render(file_name=file_name, methods=methods)
145
146
147def generate_code(request, response):
148
149 assert isinstance(request, plugin.CodeGeneratorRequest)
150 for proto_file in request.proto_file:
151 output = []
152
153 for data in traverse_methods(proto_file):
154 output.append(data)
155
156 # as a nice side-effect, generate a json file capturing the essence
157 # of the RPC method entries
158 f = response.file.add()
159 f.name = proto_file.name + '.json'
160 f.content = dumps(output, indent=4)
161
162 # generate the real Python code file
163 f = response.file.add()
164 assert proto_file.name.endswith('.proto')
165 f.name = proto_file.name.replace('.proto', '_gw.py')
166 f.content = generate_gw_code(proto_file.name, output)
167
168
169if __name__ == '__main__':
170
171 if len(sys.argv) >= 2:
172 # read input from file, to allow troubleshooting
173 with open(sys.argv[1], 'r') as f:
174 data = f.read()
175 else:
176 # read input from stdin
177 data = sys.stdin.read()
178
179 # parse request
180 request = plugin.CodeGeneratorRequest()
181 request.ParseFromString(data)
182
183 # create response object
184 response = plugin.CodeGeneratorResponse()
185
186 # generate the output and the response
187 generate_code(request, response)
188
189 # serialize the response
190 output = response.SerializeToString()
191
192 # write response to stdout
193 sys.stdout.write(output)