Skip to content

Commit 74236ea

Browse files
committed
Catch client closing the connection
1 parent 3604a93 commit 74236ea

File tree

1 file changed

+20
-10
lines changed

1 file changed

+20
-10
lines changed

jupyter_archive/handlers.py

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,11 @@
99
from notebook.utils import url2path
1010

1111

12+
# The delay in ms at which we send the chunk of data
13+
# to the client.
14+
ARCHIVE_DOWNLOAD_FLUSH_DELAY = 100
15+
16+
1217
class ArchiveStream():
1318
def __init__(self, handler):
1419
self.handler = handler
@@ -92,23 +97,20 @@ def get(self, archive_path, include_body=False):
9297
self.set_header('content-disposition',
9398
'attachment; filename={}'.format(archive_filename))
9499

95-
self.flush_cb = ioloop.PeriodicCallback(self.flush, 100)
100+
self.canceled = False
101+
self.flush_cb = ioloop.PeriodicCallback(self.flush, ARCHIVE_DOWNLOAD_FLUSH_DELAY)
96102
self.flush_cb.start()
97103

98104
args = (archive_path, archive_format, archive_token)
99-
try:
100-
yield ioloop.IOLoop.current().run_in_executor(None, self.archive_and_download, *args)
105+
yield ioloop.IOLoop.current().run_in_executor(None, self.archive_and_download, *args)
101106

102-
## FIXME: Catching the exception does not work.
103-
except Exception:
107+
if self.canceled:
104108
self.log.info('Download canceled.')
105-
106109
else:
107110
self.flush()
108-
self.set_cookie("archiveToken", archive_token)
109-
110111
self.log.info('Finished downloading {}.'.format(archive_filename))
111112

113+
self.set_cookie("archiveToken", archive_token)
112114
self.flush_cb.stop()
113115
self.finish()
114116

@@ -119,8 +121,16 @@ def archive_and_download(self, archive_path, archive_format, archive_token):
119121
for root, _, files in os.walk(archive_path):
120122
for file_ in files:
121123
file_name = os.path.join(root, file_)
122-
self.log.debug("{}\n".format(file_name))
123-
archive.add(file_name, os.path.join(root[prefix:], file_))
124+
if not self.canceled:
125+
self.log.debug("{}\n".format(file_name))
126+
archive.add(file_name, os.path.join(root[prefix:], file_))
127+
else:
128+
break
129+
130+
def on_connection_close(self):
131+
super().on_connection_close()
132+
self.canceled = True
133+
self.flush_cb.stop()
124134

125135

126136
class ExtractArchiveHandler(IPythonHandler):

0 commit comments

Comments
 (0)