fix(koa): support server-sent events in Koa (#782)

graphile · Jun 22, 2018 · 54c47be · 54c47be
1 parent 0e52908
commit 54c47be
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 17 deletions.
diff --git a/src/postgraphile/http/createPostGraphileHttpRequestHandler.js b/src/postgraphile/http/createPostGraphileHttpRequestHandler.js
@@ -462,7 +462,8 @@ export default function createPostGraphileHttpRequestHandler(options) {
       if (pathname === graphiqlRoute) {
         // If we are developing PostGraphile, instead just redirect.
         if (POSTGRAPHILE_ENV === 'development') {
-          res.writeHead(302, { Location: 'http://localhost:5783' })
+          res.statusCode = 302
+          res.setHeader('Location', 'http://localhost:5783')
           res.end()
           return
         }
@@ -766,6 +767,9 @@ export default function createPostGraphileHttpRequestHandler(options) {
 
       // Hack the req object so we can get back to ctx
       ctx.req._koaCtx = ctx
+      ctx.res.writeHead = () => {
+        throw new Error('res.writeHead not supported in Koa environment')
+      }
       ctx.res.send = () => {
         throw new Error('res.send not supported in Koa environment')
       }

diff --git a/src/postgraphile/http/setupServerSentEvents.js b/src/postgraphile/http/setupServerSentEvents.js
@@ -1,4 +1,6 @@
-export default function setupServerSentEvents (req, res, options) {
+import { PassThrough } from 'stream'
+
+export default function setupServerSentEvents(req, res, options) {
   const { _emitter } = options
 
   // Making sure these options are set.
@@ -7,18 +9,26 @@ export default function setupServerSentEvents (req, res, options) {
   req.socket.setKeepAlive(true)
 
   // Set headers for Server-Sent Events.
-  res.writeHead(200, {
-    'Content-Type': 'text/event-stream',
-    'Cache-Control': 'no-cache',
-    'Connection': 'keep-alive',
-  })
+  res.statusCode = 200
+  res.setHeader('Content-Type', 'text/event-stream')
+  res.setHeader('Cache-Control', 'no-cache')
+  res.setHeader('Connection', 'keep-alive')
+  const isKoa = !!req._koaCtx
+  const stream = isKoa ? new PassThrough() : null
+  if (isKoa) {
+    req._koaCtx.response.body = stream
+  }
 
   const sse = str => {
-    res.write(str)
-
-    // support running within the compression middleware.
-    // https://github.com/expressjs/compression#server-sent-events
-    if (res.flushHeaders) res.flushHeaders()
+    if (isKoa) {
+      stream.write(str)
+    } else {
+      res.write(str)
+
+      // support running within the compression middleware.
+      // https://github.com/expressjs/compression#server-sent-events
+      if (res.flushHeaders) res.flushHeaders()
+    }
   }
 
   // Notify client that connection is open.
@@ -27,12 +37,18 @@ export default function setupServerSentEvents (req, res, options) {
   // Setup listeners.
   const schemaChangedCb = () => sse('event: change\ndata: schema\n\n')
 
-  if (options.watchPg)
-    _emitter.on('schemas:changed', schemaChangedCb)
+  if (options.watchPg) _emitter.on('schemas:changed', schemaChangedCb)
 
   // Clean up when connection closes.
-  req.on('close', () => {
-    res.end()
+  const cleanup = () => {
+    if (stream) {
+      stream.end()
+    } else {
+      res.end()
+    }
     _emitter.removeListener('schemas:changed', schemaChangedCb)
-  })
+  }
+  req.on('close', cleanup)
+  req.on('finish', cleanup)
+  req.on('error', cleanup)
 }