Skip to content

Commit

Permalink
fix pre-commit
Browse files Browse the repository at this point in the history
  • Loading branch information
ganisback committed Aug 17, 2024
1 parent 3472e70 commit ac7ecff
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 26 deletions.
56 changes: 34 additions & 22 deletions jupyter_server_proxy/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@
Some original inspiration from https://github.com/senko/tornado-proxy
"""

import os, json, re
import os
import re
import socket
from asyncio import Lock
from copy import copy
Expand Down Expand Up @@ -305,7 +306,7 @@ def _build_proxy_request(self, host, port, proxied_path, body, **extra_opts):
decompress_response=False,
headers=headers,
**self.proxy_request_options(),
**extra_opts,
**extra_opts,
)
return req

Expand Down Expand Up @@ -374,12 +375,12 @@ async def proxy(self, host, port, proxied_path):
else:
client = httpclient.AsyncHTTPClient(force_instance=True)
# check if the request is stream request
accept_header = self.request.headers.get('Accept')
if accept_header == 'text/event-stream':
accept_header = self.request.headers.get("Accept")
if accept_header == "text/event-stream":
return await self._proxy_progressive(host, port, proxied_path, body, client)
else:
return await self._proxy_buffered(host, port, proxied_path, body, client)

async def _proxy_progressive(self, host, port, proxied_path, body, client):
# Proxy in progressive flush mode, whenever chunks are received. Potentially slower but get results quicker for voila
# Set up handlers so we can progressively flush result
Expand All @@ -388,15 +389,19 @@ async def _proxy_progressive(self, host, port, proxied_path, body, client):

def dump_headers(headers_raw):
for line in headers_raw:
r = re.match('^([a-zA-Z0-9\-_]+)\s*\:\s*([^\r\n]+)[\r\n]*$', line)
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'):
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)
self.set_header(k, v)
else:
r = re.match('^HTTP[^\s]* ([0-9]+)', line)
r = re.match(r"^HTTP[^\s]* ([0-9]+)", line)
if r:
status_code = r.group(1)
self.set_status(int(status_code))
Expand All @@ -412,20 +417,27 @@ def streaming_callback(chunk):
# record activity at start and end of requests
self._record_activity()
# 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
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)

req = self._build_proxy_request(
host,
port,
proxied_path,
body,
streaming_callback=streaming_callback,
header_callback=header_callback,
)

# no timeout for stream api
req.request_timeout = 7200
req.connect_timeout = 600

try:
response = await client.fetch(req, raise_error=False)
except httpclient.HTTPError as err:
Expand All @@ -442,15 +454,16 @@ def streaming_callback(chunk):
self.set_status(500)
self.write(str(response.error))
else:
self.set_status(response.code, response.reason) # Should already have been set
self.set_status(
response.code, response.reason
) # Should already have been set

dump_headers(headers_raw) # Should already have been emptied
dump_headers(headers_raw) # Should already have been emptied

if response.body: # Likewise, should already be chunked out and flushed
if response.body: # Likewise, should already be chunked out and flushed
self.write(response.body)

async def _proxy_buffered(self, host, port, proxied_path, body, client):

req = self._build_proxy_request(host, port, proxied_path, body)

self.log.debug(f"Proxying request to {req.url}")
Expand Down Expand Up @@ -533,7 +546,6 @@ 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.
Expand Down
6 changes: 3 additions & 3 deletions tests/resources/jupyter_server_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,10 @@ def cats_only(response, path):
response.code = 403
response.body = b"dogs not allowed"


def my_env():
return {
"MYVAR": "String with escaped {{var}}"
}
return {"MYVAR": "String with escaped {{var}}"}


c.ServerProxy.servers = {
"python-http": {
Expand Down
2 changes: 1 addition & 1 deletion tests/test_proxies.py
Original file line number Diff line number Diff line change
Expand Up @@ -375,7 +375,7 @@ def streaming_cb(data):
)
assert times_called == limit
assert all([0.45 < t < 3.0 for t in stream_read_intervals])
assert stream_data == [b'data: 0\n\n', b'data: 1\n\n', b'data: 2\n\n']
assert stream_data == [b"data: 0\n\n", b"data: 1\n\n", b"data: 2\n\n"]


async def test_server_proxy_websocket_messages(
Expand Down

0 comments on commit ac7ecff

Please sign in to comment.