-
Notifications
You must be signed in to change notification settings - Fork 21
/
station.py
225 lines (207 loc) · 9.69 KB
/
station.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
import json
import graphql
from graphql import parse as parse_graphql
from graphql import validate as validate_graphql
from jsonschema import validate
import google.protobuf.json_format as protobuf_json_format
import fastavro
from memphis.exceptions import MemphisError, MemphisSchemaError
from memphis.utils import get_internal_name
class Station:
def __init__(self, connection, name: str):
self.connection = connection
self.name = name.lower()
self.internal_station_name = get_internal_name(self.name)
async def validate_msg(self, message):
if self.connection.schema_updates_data[self.internal_station_name] != {}:
schema_type = self.connection.schema_updates_data[
self.internal_station_name
]["type"]
if schema_type == "protobuf":
message = self.validate_protobuf(message)
return message
if schema_type == "json":
message = self.validate_json_schema(message)
return message
if schema_type == "graphql":
message = self.validate_graphql(message)
return message
if schema_type == "avro":
message = self.validate_avro_schema(message)
return message
if hasattr(message, "SerializeToString"):
msg_to_send = message.SerializeToString()
return msg_to_send
elif isinstance(message, str):
message = message.encode("utf-8")
return message
elif isinstance(message, graphql.language.ast.DocumentNode):
msg = message
message = str(msg.loc.source.body)
message = message.encode("utf-8")
return message
elif hasattr(message, "SerializeToString"):
msg_to_send = message.SerializeToString()
return msg_to_send
elif not isinstance(message, bytearray) and not isinstance(message, dict):
raise MemphisSchemaError("Unsupported message type")
else:
if isinstance(message, dict):
message = bytearray(json.dumps(message).encode("utf-8"))
return message
def validate_protobuf(self, message):
proto_msg = self.connection.proto_msgs[self.internal_station_name]
msg_to_send = ""
try:
if isinstance(message, bytearray):
msg_to_send = bytes(message)
try:
proto_msg.ParseFromString(msg_to_send)
proto_msg.SerializeToString()
msg_to_send = msg_to_send.decode("utf-8")
except Exception as e:
if "parsing message" in str(e):
e = "Invalid message format, expecting protobuf"
raise MemphisSchemaError(str(e))
return message
if hasattr(message, "SerializeToString"):
msg_to_send = message.SerializeToString()
proto_msg.ParseFromString(msg_to_send)
proto_msg.SerializeToString()
try:
proto_msg.ParseFromString(msg_to_send)
proto_msg.SerializeToString()
except Exception as e:
if "parsing message" in str(e):
e = "Error parsing protobuf message"
raise MemphisSchemaError(str(e))
return msg_to_send
elif isinstance(message, dict):
try:
protobuf_json_format.ParseDict(message, proto_msg)
msg_to_send = proto_msg.SerializeToString()
return msg_to_send
except Exception as e:
raise MemphisSchemaError(str(e))
else:
raise MemphisSchemaError("Unsupported message type")
except Exception as e:
raise MemphisSchemaError("Schema validation has failed: " + str(e))
def validate_json_schema(self, message):
try:
if isinstance(message, bytearray):
try:
message_obj = json.loads(message)
except Exception as e:
raise Exception("Expecting Json format: " + str(e))
elif isinstance(message, dict):
message_obj = message
message = bytearray(json.dumps(message_obj).encode("utf-8"))
else:
raise Exception("Unsupported message type")
validate(
instance=message_obj,
schema=self.connection.json_schemas[self.internal_station_name],
)
return message
except Exception as e:
raise MemphisSchemaError("Schema validation has failed: " + str(e))
def validate_graphql(self, message):
try:
if isinstance(message, bytearray):
msg = message.decode("utf-8")
msg = parse_graphql(msg)
elif isinstance(message, str):
msg = parse_graphql(message)
message = message.encode("utf-8")
elif isinstance(message, graphql.language.ast.DocumentNode):
msg = message
message = str(msg.loc.source.body)
message = message.encode("utf-8")
else:
raise Exception("Unsupported message type")
validate_res = validate_graphql(
schema=self.connection.graphql_schemas[self.internal_station_name],
document_ast=msg,
)
if len(validate_res) > 0:
raise Exception(
"Schema validation has failed: " + str(validate_res))
return message
except Exception as e:
if "Syntax Error" in str(e):
e = "Invalid message format, expected GraphQL"
raise MemphisSchemaError("Schema validation has failed: " + str(e))
def validate_avro_schema(self, message):
try:
if isinstance(message, bytearray):
try:
message_obj = json.loads(message)
except Exception as e:
raise Exception("Expecting Avro format: " + str(e))
elif isinstance(message, dict):
message_obj = message
message = bytearray(json.dumps(message_obj).encode("utf-8"))
else:
raise Exception("Unsupported message type")
fastavro.validate(
message_obj,
self.connection.avro_schemas[self.internal_station_name],
)
return message
except fastavro.validation.ValidationError as e:
raise MemphisSchemaError("Schema validation has failed: " + str(e))
async def destroy(self, timeout_retries=5):
"""Destroy the station."""
try:
name_req = {"station_name": self.name, "username": self.connection.username}
station_name = json.dumps(name_req, indent=2).encode("utf-8")
# pylint: disable=protected-access
res = await self.connection._request(
"$memphis_station_destructions", station_name, 5, timeout_retries
)
# pylint: enable=protected-access
error = res.data.decode("utf-8")
if error != "" and not "not exist" in error:
raise MemphisError(error)
internal_station_name = get_internal_name(self.name)
sub = self.connection.schema_updates_subs.get(internal_station_name)
task = self.connection.schema_tasks.get(internal_station_name)
if internal_station_name in self.connection.schema_updates_data:
del self.connection.schema_updates_data[internal_station_name]
if internal_station_name in self.connection.schema_updates_subs:
del self.connection.schema_updates_subs[internal_station_name]
if internal_station_name in self.connection.clients_per_station:
del self.connection.clients_per_station[internal_station_name]
if internal_station_name in self.connection.schema_tasks:
del self.connection.schema_tasks[internal_station_name]
if task is not None:
task.cancel()
if sub is not None:
await sub.unsubscribe()
if internal_station_name in self.connection.functions_clients_per_station:
del self.connection.functions_clients_per_station[internal_station_name]
if internal_station_name in self.connection.functions_updates_data:
del self.connection.functions_updates_data[internal_station_name]
if internal_station_name in self.connection.functions_updates_subs:
function_sub = self.connection.functions_updates_subs.get(internal_station_name)
if function_sub is not None:
await function_sub.unsubscribe()
del self.connection.functions_updates_subs[internal_station_name]
if internal_station_name in self.connection.functions_tasks:
task = self.connection.functions_tasks.get(internal_station_name)
if task is not None:
task.cancel()
del self.connection.functions_tasks[internal_station_name]
self.connection.producers_map = {
k: v
for k, v in self.connection.producers_map.items()
if self.name not in k
}
self.connection.consumers_map = {
k: v
for k, v in self.connection.consumers_map.items()
if self.name not in k
}
except Exception as e:
raise MemphisError(str(e)) from e