# subscription_client.py from base64 import b64encode, decode from datetime import datetime from uuid import uuid4 import websocket import threading import json from functools import partial # Constants Copied from AppSync API 'Settings' API_URL = "https://rqdq6pkcijbrpdqi5egplq4rvu.appsync-api.eu-west-1.amazonaws.com/graphql" API_KEY = "da2-yjil26sgsfajvhxx3lqdbosv4q" # GraphQL subscription Registration object GQL_SUBSCRIPTION = json.dumps({ 'query': 'subscription test { changedUser { __typename UID } }', 'variables': {} }) APPSYNC_REALTIME_HEADERS = { 'accept': 'application/json, text/javascript', 'content-encoding': 'amz-1.0', 'content-type': 'application/json; charset=UTF-8' } # Discovered values from the AppSync endpoint (API_URL) WSS_URL = API_URL.replace('https', 'wss').replace('appsync-api', 'appsync-realtime-api') HOST = API_URL.replace('https://', '').replace('/graphql', '') # Set up Timeout Globals timeout_timer = None timeout_interval = 10 # Calculate UTC time in ISO format (AWS Friendly): YYYY-MM-DDTHH:mm:ssZ def header_time(): return datetime.utcnow().isoformat(sep='T', timespec='seconds') + 'Z' def header_encode(header_obj): """ Encode Using Base 64""" return b64encode(json.dumps(header_obj).encode('utf-8')).decode('utf-8') # reset the keep alive timeout daemon thread def reset_timer(ws): global timeout_timer global timeout_interval if timeout_timer: timeout_timer.cancel() timeout_timer = threading.Timer(timeout_interval, lambda: ws.close()) timeout_timer.daemon = True timeout_timer.start() class ClientCallback: def __init__(self, on_open=None, on_message=None, on_error=None, on_close=None): self.on_open = on_open self.on_message = on_message self.on_error = on_error self.on_close = on_close class Subscription: def __init__(self, wss_url, url, host, gql_query, auth_token=None): self.wss_url = wss_url self.url = url self.host = host self.auth_header = APPSYNC_REALTIME_HEADERS self.req_header = { 'authorization': { 'Authorization': auth_token, 'host': host, 'x-amz-user-agent': 'fansifter-py/0.0' } } self.api_header = { 'host': host, 'authorization': auth_token } # 'x-api-key': api_id self.gql_query = gql_query self.SUB_ID = str(uuid4()) self.socket_app = None self.client_callback = None self.trace = False def sub(self, trace=False, on_open=None, on_message=None, on_error=None, on_close=None, auth_token=None, run_forever=True): if trace: websocket.enableTrace(True) self.trace = trace if self.trace: print('Connecting to: ' + self.wss_url) if auth_token: self.req_header = { 'authorization': { 'Authorization': auth_token, 'host': self.host, 'x-amz-user-agent': 'fansifter-py/0.0' } } self.api_header = { 'host': self.host, 'authorization': auth_token } # payload of "e30=" = b64encode('{}'). self.client_callback = ClientCallback(on_open=on_open, on_message=on_message, on_error=on_error, on_close=on_close) self.socket_app = websocket.WebSocketApp(self.wss_url + '?header=' + header_encode(self.api_header) + '&payload=e30=', header=self.auth_header, subprotocols=['graphql-ws'], on_open=partial(self.on_open), on_message=partial(self.on_message), on_error=partial(self.on_error), on_close=partial(self.on_close), ) if run_forever: self.socket_app.run_forever() def run_forever(self): if self.socket_app: self.socket_app.run_forever() def stop(self): if self.socket_app is not None: deregister = { 'type': 'stop', 'id': self.SUB_ID } end_sub = json.dumps(deregister) if self.trace: print('>> ' + end_sub) self.socket_app.send(end_sub) def on_message(self, ws, message): global timeout_timer global timeout_interval if self.trace: print('### message ###') print('<< ' + message) message_object = json.loads(message) message_type = message_object['type'] if message_type == 'data': pass elif message_type == 'ka': reset_timer(ws) elif message_type == 'connection_ack': timeout_interval = int(json.dumps(message_object['payload']['connectionTimeoutMs'])) register = { 'id': self.SUB_ID, 'payload': { 'data': self.gql_query, 'extensions': self.req_header }, 'type': 'start' } start_sub = json.dumps(register) if self.trace: print('>> ' + start_sub) ws.send(start_sub) elif message_type == "complete": if self.trace: print('### complete ###') self.socket_app.close() elif self.trace and message_type == 'error': print('Error from AppSync: ' + json.dumps(message_object['payload'])) if self.client_callback.on_message: self.client_callback.on_message(message_type, message_object) def on_error(self, ws, error): if self.trace: print('### error ###') print(error) if self.client_callback.on_error: self.client_callback.on_error(error) def on_close(self, ws): if self.trace: print('### closed ###') if self.client_callback.on_close: self.client_callback.on_close() def on_open(self, ws): if self.trace: print('### opened ###') init = { 'type': 'connection_init' } init_conn = json.dumps(init) if self.trace: print('>> ' + init_conn) ws.send(init_conn) if self.client_callback.on_open: self.client_callback.on_open() # if __name__ == '__main__': # # Uncomment to see socket bytestreams # # )websocket.enableTrace(True # # # Set up the connection URL, which includes the Authentication Header # # and a payload of '{}'. All info is base 64 encoded # connection_url = WSS_URL + '?header=' + header_encode(api_header) + '&payload=e30=' # # # Create the websocket connection to AppSync's real-time endpoint # # also defines callback functions for websocket events # # NOTE: The connection requires a subprotocol 'graphql-ws' # print('Connecting to: ' + connection_url) # # socket = websocket.WebSocketApp(connection_url, # subprotocols=['graphql-ws'], # on_open=on_open, # on_message=on_message, # on_error=on_error, # on_close=on_close, ) # # socket.run_forever()