support stream

Author: ganisback

jupyter_server_proxy/handlers.py

@@ -4,7 +4,7 @@
 Some original inspiration from https://github.com/senko/tornado-proxy
 """
 
-import os
+import os, json, re
 import socket
 from asyncio import Lock
 from copy import copy
@@ -287,7 +287,7 @@ def get_client_uri(self, protocol, host, port, proxied_path):
 
         return client_uri
 
-    def _build_proxy_request(self, host, port, proxied_path, body):
+    def _build_proxy_request(self, host, port, proxied_path, body, **extra_opts):
         headers = self.proxy_request_headers()
 
         client_uri = self.get_client_uri("http", host, port, proxied_path)
@@ -307,6 +307,7 @@ def _build_proxy_request(self, host, port, proxied_path, body):
             decompress_response=False,
             headers=headers,
             **self.proxy_request_options(),
+             **extra_opts,
         )
         return req
 
@@ -365,7 +366,82 @@ async def proxy(self, host, port, proxied_path):
                 body = b""
             else:
                 body = None
+        accept_type = self.request.headers.get('Accept')
+        if accept_type == 'text/event-stream':
+            return await self._proxy_progressive(host, port, proxied_path, body)
+        else:
+            return await self._proxy_normal(host, port, proxied_path, body)
+    
+    async def _proxy_progressive(self, host, port, proxied_path, body):
+        # Proxy in progressive flush mode, whenever chunks are received. Potentially slower but get results quicker for voila
+
+        client = httpclient.AsyncHTTPClient()
+
+        # Set up handlers so we can progressively flush result
+
+        headers_raw = []
+
+        def dump_headers(headers_raw):
+            for line in headers_raw:
+                r = re.match('^([a-zA-Z0-9\-_]+)\s*\:\s*([^\r\n]+)[\r\n]*$', line)
+                if r:
+                    k,v = r.groups([1,2])
+                    if k not in ('Content-Length', 'Transfer-Encoding',
+                                  'Content-Encoding', 'Connection'):
+                        # some header appear multiple times, eg 'Set-Cookie'
+                        self.set_header(k,v)
+                else:
+                    r = re.match('^HTTP[^\s]* ([0-9]+)', line)
+                    if r:
+                        status_code = r.group(1)
+                        self.set_status(int(status_code))
+            headers_raw.clear()
+
+        # clear tornado default header
+        self._headers = httputil.HTTPHeaders()
+
+        def header_callback(line):
+            headers_raw.append(line)
+
+        def streaming_callback(chunk):
+            # Do this here, not in header_callback so we can be sure headers are out of the way first
+            dump_headers(headers_raw) # array will be empty if this was already called before
+            self.write(chunk)
+            self.flush()
+
+        # Now make the request
+
+        req = self._build_proxy_request(host, port, proxied_path, body, 
+                    streaming_callback=streaming_callback, 
+                    header_callback=header_callback)
+
+        try:
+            response = await client.fetch(req, raise_error=False)
+        except httpclient.HTTPError as err:
+            if err.code == 599:
+                self._record_activity()
+                self.set_status(599)
+                self.write(str(err))
+                return
+            else:
+                raise
+
+        # record activity at start and end of requests
+        self._record_activity()
+
+        # For all non http errors...
+        if response.error and type(response.error) is not httpclient.HTTPError:
+            self.set_status(500)
+            self.write(str(response.error))
+        else:
+            self.set_status(response.code, response.reason) # Should already have been set
+
+            dump_headers(headers_raw) # Should already have been emptied
 
+            if response.body: # Likewise, should already be chunked out and flushed
+                self.write(response.body)
+
+    async def _proxy_normal(self, host, port, proxied_path, body):
         if self.unix_socket is not None:
             # Port points to a Unix domain socket
             self.log.debug("Making client for Unix socket %r", self.unix_socket)
@@ -458,6 +534,7 @@ def rewrite_pe(rewritable_response: RewritableResponse):
             if rewritten_response.body:
                 self.write(rewritten_response.body)
 
+
     async def proxy_open(self, host, port, proxied_path=""):
         """
         Called when a client opens a websocket connection.