From 3294f77ca5adfd22a494d7b1fe65d6223663dcaa Mon Sep 17 00:00:00 2001
From: Ben Reinhart <benjreinhart@gmail.com>
Date: Thu, 15 Apr 2021 09:42:01 -0700
Subject: [PATCH] feat: Add health endpoint to WebSocket server (#14110)

---
 superset-websocket/README.md          | 20 +++++++++-
 superset-websocket/spec/index.test.ts | 54 +++++++++++++++++++++++++++
 superset-websocket/src/index.ts       | 22 +++++++++++
 3 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/superset-websocket/README.md b/superset-websocket/README.md
index b5cd5579f..aa6a9f201 100644
--- a/superset-websocket/README.md
+++ b/superset-websocket/README.md
@@ -81,6 +81,8 @@ GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL = "ws://<host>:<port>/"
 
 Note that the WebSocket server must be run on the same hostname (different port) for cookies to be shared between the Flask app and the WebSocket server.
 
+Note also that `localhost` and `127.0.0.1` are not considered the same host. For example, if you're pointing your browser to `localhost:<port>` for Superset, then the WebSocket url will need to be configured as `localhost:<port>`.
+
 The following config values must contain the same values in both the Flask app config and `config.json`:
 ```
 GLOBAL_ASYNC_QUERIES_REDIS_CONFIG
@@ -103,4 +105,20 @@ Running in production:
 npm run build && npm start
 ```
 
-*TODO: containerization*
+## Health check
+
+The WebSocket server supports health checks via one of:
+
+```
+GET /health
+```
+
+OR
+
+```
+HEAD /health
+```
+
+## Containerization
+
+*TODO: containerize websocket server*
diff --git a/superset-websocket/spec/index.test.ts b/superset-websocket/spec/index.test.ts
index 5e703c4dd..c92a553a2 100644
--- a/superset-websocket/spec/index.test.ts
+++ b/superset-websocket/spec/index.test.ts
@@ -63,6 +63,60 @@ describe('server', () => {
     server.resetState();
   });
 
+  describe('HTTP requests', () => {
+    test('services health checks', () => {
+      const endMock = jest.fn();
+      const writeHeadMock = jest.fn();
+
+      const request = {
+        url: '/health',
+        method: 'GET',
+        headers: {
+          host: 'example.com',
+        },
+      };
+
+      const response = {
+        writeHead: writeHeadMock,
+        end: endMock,
+      };
+
+      server.httpRequest(request as any, response as any);
+
+      expect(writeHeadMock).toBeCalledTimes(1);
+      expect(writeHeadMock).toHaveBeenLastCalledWith(200);
+
+      expect(endMock).toBeCalledTimes(1);
+      expect(endMock).toHaveBeenLastCalledWith('OK');
+    });
+
+    test('reponds with a 404 otherwise', () => {
+      const endMock = jest.fn();
+      const writeHeadMock = jest.fn();
+
+      const request = {
+        url: '/unsupported',
+        method: 'GET',
+        headers: {
+          host: 'example.com',
+        },
+      };
+
+      const response = {
+        writeHead: writeHeadMock,
+        end: endMock,
+      };
+
+      server.httpRequest(request as any, response as any);
+
+      expect(writeHeadMock).toBeCalledTimes(1);
+      expect(writeHeadMock).toHaveBeenLastCalledWith(404);
+
+      expect(endMock).toBeCalledTimes(1);
+      expect(endMock).toHaveBeenLastCalledWith('Not Found');
+    });
+  });
+
   describe('incrementId', () => {
     test('it increments a valid Redis stream ID', () => {
       expect(server.incrementId('1607477697866-0')).toEqual('1607477697866-1');
diff --git a/superset-websocket/src/index.ts b/superset-websocket/src/index.ts
index 9464295c8..14ae5f829 100644
--- a/superset-websocket/src/index.ts
+++ b/superset-websocket/src/index.ts
@@ -349,6 +349,27 @@ export const wsConnection = (ws: WebSocket, request: http.IncomingMessage) => {
   });
 };
 
+/**
+ * HTTP `request` event handler, called via httpServer
+ */
+export const httpRequest = (
+  request: http.IncomingMessage,
+  response: http.ServerResponse,
+) => {
+  const rawUrl = request.url as string;
+  const method = request.method as string;
+  const headers = request.headers || {};
+  const url = new URL(rawUrl as string, `http://${headers.host}`);
+  if (url.pathname === '/health' && ['GET', 'HEAD'].includes(method)) {
+    response.writeHead(200);
+    response.end('OK');
+  } else {
+    logger.info(`Received unexpected request: ${method} ${rawUrl}`);
+    response.writeHead(404);
+    response.end('Not Found');
+  }
+};
+
 /**
  * HTTP `upgrade` event handler, called via httpServer
  */
@@ -439,6 +460,7 @@ export const cleanChannel = (channel: string) => {
 if (startServer) {
   // init server event listeners
   wss.on('connection', wsConnection);
+  httpServer.on('request', httpRequest);
   httpServer.on('upgrade', httpUpgrade);
   httpServer.listen(opts.port);
   logger.info(`Server started on port ${opts.port}`);