diff --git a/zappa/cli.py b/zappa/cli.py index 1a07583ea..f54e224a0 100644 --- a/zappa/cli.py +++ b/zappa/cli.py @@ -111,6 +111,7 @@ class ZappaCLI(object): environment_variables = None authorizer = None aws_kms_key_arn = '' + web_socket_handler = None cognito_authorizer_path = '/' @@ -1820,6 +1821,9 @@ def load_settings(self, settings_file=None, session=None): self.aws_kms_key_arn = self.stage_config.get('aws_kms_key_arn', '') self.cognito_authorizer_path = self.stage_config.get('cognito_authorizer_path', '/') + # custom zappa settings key for web socket + self.web_socket_handler = self.stage_config.get('web_socket_handler') + desired_role_name = self.lambda_name + "-ZappaLambdaExecutionRole" self.zappa = Zappa( boto_session=session, profile_name=self.profile_name, @@ -2090,6 +2094,12 @@ def create_package(self, output=None): if authorizer_function: settings_s += "AUTHORIZER_FUNCTION='{0!s}'\n".format(authorizer_function) + # web socket function + if self.web_socket_handler: + settings_s += "WEB_SOCKET_HANDLER='{0!s}'\n".format( + self.web_socket_handler) + else: + settings_s += "WEB_SOCKET_HANDLER=None\n" # Copy our Django app into root of our package. # It doesn't work otherwise. diff --git a/zappa/handler.py b/zappa/handler.py index a1881e53f..5c564a357 100644 --- a/zappa/handler.py +++ b/zappa/handler.py @@ -425,6 +425,20 @@ def handler(self, event, context): logger.error("Cannot find a function to process the authorization request.") raise Exception('Unauthorized') + # this is a web socket event + elif event.get('requestContext') and \ + event['requestContext'].get('connectionId'): + whole_function = self.settings.WEB_SOCKET_HANDLER + if whole_function: + app_function = self.import_module_and_get_function( + whole_function) + result = self.run_function(app_function, event, context) + return result + else: + logger.error( + "Cannot find a function to process the web socket request.") + raise Exception('Web socket request handler not found') + # Normal web app flow try: # Timing