Repository: streamlit/streamlit

Issue:
File Uploader: Consolidate reruns for multiple files
Solution:
Every time a file uploads, rerun only if there are no other incomplete uploads (in the same widget)

Alternatives discussed:
- Parallel upload == 1 rerun. Sequential = N
- (super fancy) Like today, but throttled. Only rerun if no upload in progress after.
- Add fancy parameters
    - Customizable v1. Introduce a parameter that lets you switch behaviors
    - Customizable v2.  Developer can specify number of files, and rerun only when N is hit

[More details here](https://www.notion.so/streamlit/Consolidate-reruns-for-multiple-file-uploads-97e07fc4474a452ab7cd35dfda70b9d6)


Discussion:

Resolving patch:
diff --git a/lib/streamlit/server/server.py b/lib/streamlit/server/server.py
--- a/lib/streamlit/server/server.py
+++ b/lib/streamlit/server/server.py
@@ -336,20 +336,13 @@ def _create_app(self):
                 dict(cache=self._message_cache),
             ),
             (
-                # Tornado doesn't allow body in DELETE requests.
-                # Passing lookup values in URL
                 make_url_path_regex(
                     base,
-                    "/upload_file/(?P<session_id>.*)/(?P<widget_id>.*)/(?P<file_id>[0-9]*)?",
+                    "/upload_file/?(?P<session_id>[^/]*)?/?(?P<widget_id>[^/]*)?/?(?P<file_id>[0-9]*)?",
                 ),
                 UploadFileRequestHandler,
                 dict(file_mgr=self._uploaded_file_mgr),
             ),
-            (
-                make_url_path_regex(base, "upload_file"),
-                UploadFileRequestHandler,
-                dict(file_mgr=self._uploaded_file_mgr),
-            ),
             (
                 make_url_path_regex(base, "assets/(.*)"),
                 AssetsFileHandler,
diff --git a/lib/streamlit/server/upload_file_request_handler.py b/lib/streamlit/server/upload_file_request_handler.py
--- a/lib/streamlit/server/upload_file_request_handler.py
+++ b/lib/streamlit/server/upload_file_request_handler.py
@@ -44,14 +44,14 @@ def initialize(self, file_mgr):
         self._file_mgr = file_mgr
 
     def set_default_headers(self):
-        self.set_header("Access-Control-Allow-Methods", "POST, DELETE")
+        self.set_header("Access-Control-Allow-Methods", "POST, PUT, DELETE")
         self.set_header("Access-Control-Allow-Headers", "Content-Type")
         if config.get_option("server.enableXsrfProtection"):
-            self.set_header("Access-Control-Allow-Headers", "X-Xsrftoken, Content-Type")
             self.set_header(
                 "Access-Control-Allow-Origin",
                 Report.get_url(config.get_option("browser.serverAddress")),
             )
+            self.set_header("Access-Control-Allow-Headers", "X-Xsrftoken, Content-Type")
             self.set_header("Vary", "Origin")
             self.set_header("Access-Control-Allow-Credentials", "true")
         elif routes.allow_cross_origin_requests():
@@ -96,7 +96,7 @@ def _require_arg(args, name):
         # Convert bytes to string
         return arg[0].decode("utf-8")
 
-    def post(self):
+    def post(self, **kwargs):
         args = {}  # type: Dict[str, List[bytes]]
         files = {}  # type: Dict[str, List[Any]]
 
@@ -114,7 +114,9 @@ def post(self):
             self.send_error(400, reason=str(e))
             return
 
-        LOGGER.debug(f"{len(files)} file(s) received for session {session_id} widget {widget_id}")
+        LOGGER.debug(
+            f"{len(files)} file(s) received for session {session_id} widget {widget_id}"
+        )
 
         # Create an UploadedFile object for each file.
         uploaded_files = []
@@ -132,11 +134,12 @@ def post(self):
         if len(uploaded_files) == 0:
             self.send_error(400, reason="Expected at least 1 file, but got 0")
             return
-
         replace = self.get_argument("replace", "false")
 
         update_files = (
-            self._file_mgr.replace_files if replace == "true" else self._file_mgr.add_files
+            self._file_mgr.replace_files
+            if replace == "true"
+            else self._file_mgr.add_files
         )
         update_files(
             session_id=session_id,
@@ -144,10 +147,26 @@ def post(self):
             files=uploaded_files,
         )
 
-        LOGGER.debug(f"{len(files)} file(s) uploaded for session {session_id} widget {widget_id}. replace {replace}")
+        LOGGER.debug(
+            f"{len(files)} file(s) uploaded for session {session_id} widget {widget_id}. replace {replace}"
+        )
 
         self.set_status(200)
 
+    def put(self, session_id: str, widget_id: str, **kwargs):
+        if self.request.body:
+            data = tornado.escape.json_decode(self.request.body)
+            if "totalFiles" in data.keys():
+                self._file_mgr.update_file_count(
+                    session_id=session_id,
+                    widget_id=widget_id,
+                    file_count=data["totalFiles"],
+                )
+                self.set_status(200)
+                return
+
+        self.send_error(400, reason="Nothing to update")
+
     def delete(self, session_id, widget_id, file_id):
         if session_id is None or widget_id is None or file_id is None:
             self.send_error(404)
diff --git a/lib/streamlit/uploaded_file_manager.py b/lib/streamlit/uploaded_file_manager.py
--- a/lib/streamlit/uploaded_file_manager.py
+++ b/lib/streamlit/uploaded_file_manager.py
@@ -37,6 +37,7 @@ class UploadedFileManager(object):
 
     def __init__(self):
         self._files_by_id = {}  # type: Dict[Tuple[str, str], List[UploadedFile] ]
+        self._file_counts_by_id: Dict[Tuple[str, str], int] = {}
         # Prevents concurrent access to the _files_by_id dict.
         # In remove_session_files(), we iterate over the dict's keys. It's
         # an error to mutate a dict while iterating; this lock prevents that.
@@ -51,10 +52,26 @@ def __init__(self):
             """
         )
 
-    def _on_files_updated(self, session_id):
-        self.on_files_updated.send(session_id)
-
-    def _add_files(self, session_id, widget_id, files):
+    def _on_files_updated(self, session_id: str, widget_id: str):
+        files_by_widget = session_id, widget_id
+        if files_by_widget in self._file_counts_by_id:
+            expected_file_count: int = self._file_counts_by_id[files_by_widget]
+            actual_file_count: int = (
+                len(self._files_by_id[files_by_widget])
+                if files_by_widget in self._files_by_id
+                else 0
+            )
+            if expected_file_count == actual_file_count:
+                self.on_files_updated.send(session_id)
+        else:
+            self.on_files_updated.send(session_id)
+
+    def _add_files(
+        self,
+        session_id: str,
+        widget_id: str,
+        files: List[UploadedFile],
+    ):
         """
         Add a list of files to the FileManager. Does not emit any signals
         """
@@ -65,7 +82,12 @@ def _add_files(self, session_id, widget_id, files):
                 files = file_list + files
             self._files_by_id[files_by_widget] = files
 
-    def add_files(self, session_id, widget_id, files):
+    def add_files(
+        self,
+        session_id: str,
+        widget_id: str,
+        files: List[UploadedFile],
+    ):
         """Add a list of files to the FileManager.
 
         The "on_file_added" Signal will be emitted after the list is added.
@@ -79,10 +101,11 @@ def add_files(self, session_id, widget_id, files):
         files : List[UploadedFile]
             The files to add.
         """
+
         self._add_files(session_id, widget_id, files)
-        self._on_files_updated(session_id)
+        self._on_files_updated(session_id, widget_id)
 
-    def get_files(self, session_id, widget_id):
+    def get_files(self, session_id: str, widget_id: str):
         """Return the file list with the given ID, or None if the ID doesn't exist.
 
         Parameters
@@ -96,39 +119,34 @@ def get_files(self, session_id, widget_id):
         -------
         list of UploadedFile or None
         """
-        files_id = session_id, widget_id
+        files_by_widget = session_id, widget_id
         with self._files_lock:
-            file_list = self._files_by_id.get(files_id, None)
+            file_list = self._files_by_id.get(files_by_widget, None)
         return file_list
 
-    def remove_file(self, session_id, widget_id, file_id):
-        """Remove the file list with the given ID, if it exists.
-        Parameters
-        ----------
-        session_id : str
-            The session ID of the report that owns the file.
-        widget_id : str
-            The widget ID of the FileUploader that created the file.
-        """
+    def remove_file(self, session_id: str, widget_id: str, file_id: str):
+        """Remove the file list with the given ID, if it exists."""
         files_by_widget = session_id, widget_id
         with self._files_lock:
             file_list = self._files_by_id[files_by_widget]
             self._files_by_id[files_by_widget] = [
                 file for file in file_list if file.id != file_id
             ]
-        self._on_files_updated(session_id)
+            if len(file_list) != len(self._files_by_id[files_by_widget]):
+                self._on_files_updated(session_id, widget_id)
 
-    def _remove_files(self, session_id, widget_id):
+    def _remove_files(self, session_id: str, widget_id: str):
         """Remove the file list for the provided widget in the
         provided session, if it exists.
 
         Does not emit any signals.
         """
-        files_id = session_id, widget_id
+        files_by_widget = session_id, widget_id
+        self.update_file_count(session_id, widget_id, 0)
         with self._files_lock:
-            self._files_by_id.pop(files_id, None)
+            self._files_by_id.pop(files_by_widget, None)
 
-    def remove_files(self, session_id, widget_id):
+    def remove_files(self, session_id: str, widget_id: str):
         """Remove the file list for the provided widget in the
         provided session, if it exists.
 
@@ -140,9 +158,9 @@ def remove_files(self, session_id, widget_id):
             The widget ID of the FileUploader that created the file.
         """
         self._remove_files(session_id, widget_id)
-        self._on_files_updated(session_id)
+        self._on_files_updated(session_id, widget_id)
 
-    def remove_session_files(self, session_id):
+    def remove_session_files(self, session_id: str):
         """Remove all files that belong to the given report.
 
         Parameters
@@ -159,7 +177,12 @@ def remove_session_files(self, session_id):
             if files_id[0] == session_id:
                 self.remove_files(*files_id)
 
-    def replace_files(self, session_id, widget_id, files):
+    def replace_files(
+        self,
+        session_id: str,
+        widget_id: str,
+        files: List[UploadedFile],
+    ):
         """Removes the file list for the provided widget in the
         provided session, if it exists and add the provided files
         to the widget in the session
@@ -175,4 +198,14 @@ def replace_files(self, session_id, widget_id, files):
         """
         self._remove_files(session_id, widget_id)
         self._add_files(session_id, widget_id, files)
-        self._on_files_updated(session_id)
+        self._on_files_updated(session_id, widget_id)
+
+    def update_file_count(
+        self,
+        session_id: str,
+        widget_id: str,
+        file_count: int,
+    ):
+        files_by_widget = session_id, widget_id
+        self._file_counts_by_id[files_by_widget] = file_count
+        self._on_files_updated(session_id, widget_id)


Test patch:
diff --git a/frontend/src/components/widgets/FileUploader/FileUploader.test.tsx b/frontend/src/components/widgets/FileUploader/FileUploader.test.tsx
--- a/frontend/src/components/widgets/FileUploader/FileUploader.test.tsx
+++ b/frontend/src/components/widgets/FileUploader/FileUploader.test.tsx
@@ -47,6 +47,7 @@ const FILE_TOO_LARGE: FileError = {
 
 const getProps = (elementProps: Record<string, unknown> = {}): Props => ({
   element: FileUploaderProto.create({
+    id: "id",
     type: [],
     maxUploadSizeMb: 50,
     ...elementProps,
@@ -56,7 +57,10 @@ const getProps = (elementProps: Record<string, unknown> = {}): Props => ({
   // @ts-ignore
   widgetStateManager: jest.fn(),
   // @ts-ignore
-  uploadClient: { uploadFiles: jest.fn().mockResolvedValue(undefined) },
+  uploadClient: {
+    uploadFiles: jest.fn().mockResolvedValue(undefined),
+    updateFileCount: jest.fn().mockResolvedValue(undefined),
+  },
 })
 
 describe("FileUploader widget", () => {
@@ -98,6 +102,32 @@ describe("FileUploader widget", () => {
     )
 
     expect(props.uploadClient.uploadFiles.mock.calls.length).toBe(1)
+    expect(props.uploadClient.updateFileCount.mock.calls.length).toBe(1)
+    expect(props.uploadClient.updateFileCount).toBeCalledWith(
+      props.element.id,
+      1
+    )
+  })
+
+  it("should upload multiple files", () => {
+    const props = getProps({ multipleFiles: true })
+    const wrapper = shallow(<FileUploader {...props} />)
+    const internalFileUploader = wrapper.find(FileDropzone)
+    internalFileUploader.props().onDrop(
+      [blobFile, blobFile],
+      [
+        { file: blobFile, errors: [INVALID_TYPE_ERROR, TOO_MANY_FILES] },
+        { file: blobFile, errors: [TOO_MANY_FILES] },
+        { file: blobFile, errors: [TOO_MANY_FILES] },
+      ]
+    )
+
+    expect(props.uploadClient.uploadFiles.mock.calls.length).toBe(2)
+    expect(props.uploadClient.updateFileCount.mock.calls.length).toBe(1)
+    expect(props.uploadClient.updateFileCount).toBeCalledWith(
+      props.element.id,
+      2
+    )
   })
 
   it("should change status + add file attributes when dropping a File", () => {
diff --git a/lib/tests/streamlit/server_test.py b/lib/tests/streamlit/server_test.py
--- a/lib/tests/streamlit/server_test.py
+++ b/lib/tests/streamlit/server_test.py
@@ -329,6 +329,11 @@ def test_uploaded_file_triggers_rerun(self):
             file = UploadedFile("id", "file.txt", "type", b"123")
 
             # "Upload a file" for Session1
+            self.server._uploaded_file_mgr.update_file_count(
+                session_id=session_info1.session.id,
+                widget_id="widget_id",
+                file_count=1,
+            )
             self.server._uploaded_file_mgr.add_files(
                 session_id=session_info1.session.id,
                 widget_id="widget_id",
@@ -346,6 +351,44 @@ def test_uploaded_file_triggers_rerun(self):
             session_info1.session.request_rerun.assert_called_once()
             session_info2.session.request_rerun.assert_not_called()
 
+    @tornado.testing.gen_test
+    def test_multiple_uploaded_file_triggers_one_rerun(self):
+        """Uploading a file should trigger a re-run in the associated
+        ReportSession."""
+        with self._patch_report_session():
+            yield self.start_server_loop()
+
+            # Connect twice and get associated ReportSessions
+            yield self.ws_connect()
+            yield self.ws_connect()
+            session_info1 = list(self.server._session_info_by_id.values())[0]
+            session_info2 = list(self.server._session_info_by_id.values())[1]
+
+            file = UploadedFile("id", "file.txt", "type", b"123")
+
+            # "Upload 2 files" for Session1
+            self.server._uploaded_file_mgr.update_file_count(
+                session_id=session_info1.session.id,
+                widget_id="widget_id",
+                file_count=2,
+            )
+            self.server._uploaded_file_mgr.add_files(
+                session_id=session_info1.session.id,
+                widget_id="widget_id",
+                files=[file, file],
+            )
+
+            self.assertEqual(
+                self.server._uploaded_file_mgr.get_files(
+                    session_info1.session.id, "widget_id"
+                ),
+                [file, file],
+            )
+
+            # Session1 should have a rerun request; Session2 should not
+            session_info1.session.request_rerun.assert_called_once()
+            session_info2.session.request_rerun.assert_not_called()
+
     @tornado.testing.gen_test
     def test_orphaned_upload_file_deletion(self):
         """An uploaded file with no associated ReportSession should be
diff --git a/lib/tests/streamlit/upload_file_request_handler_test.py b/lib/tests/streamlit/upload_file_request_handler_test.py
--- a/lib/tests/streamlit/upload_file_request_handler_test.py
+++ b/lib/tests/streamlit/upload_file_request_handler_test.py
@@ -75,6 +75,7 @@ def test_upload_one_file(self):
             file.name: file,
             "sessionId": (None, "fooReport"),
             "widgetId": (None, "barWidget"),
+            "totalFiles": (None, "1"),
         }
         response = self._upload_files(params)
         self.assertEqual(200, response.code)
@@ -97,6 +98,7 @@ def test_upload_multiple_files(self):
             file3.name: file3,
             "sessionId": (None, "fooReport"),
             "widgetId": (None, "barWidget"),
+            "totalFiles": (None, "1"),
         }
         response = self._upload_files(params)
         self.assertEqual(200, response.code)
@@ -116,6 +118,7 @@ def test_upload_missing_params(self):
             "image.png": ("image.png", b"1234"),
             "sessionId": (None, "fooReport"),
             # "widgetId": (None, 'barWidget'),
+            "totalFiles": (None, "1"),
         }
 
         response = self._upload_files(params)
@@ -128,6 +131,7 @@ def test_upload_missing_file(self):
             # "image.png": ("image.png", b"1234"),
             "sessionId": (None, "fooReport"),
             "widgetId": (None, "barWidget"),
+            "totalFiles": (None, "1"),
         }
         response = self._upload_files(params)
         self.assertEqual(400, response.code)
diff --git a/lib/tests/streamlit/uploaded_file_manager_test.py b/lib/tests/streamlit/uploaded_file_manager_test.py
--- a/lib/tests/streamlit/uploaded_file_manager_test.py
+++ b/lib/tests/streamlit/uploaded_file_manager_test.py
@@ -35,11 +35,13 @@ def _on_files_updated(self, file_list, **kwargs):
     def test_add_file(self):
         self.assertIsNone(self.mgr.get_files("non-report", "non-widget"))
 
+        self.mgr.update_file_count("session", "widget", 1)
         self.mgr.add_files("session", "widget", [file1])
         self.assertEqual([file1], self.mgr.get_files("session", "widget"))
         self.assertEqual(len(self.filemgr_events), 1)
 
         # Add another file with the same ID
+        self.mgr.update_file_count("session", "widget", 2)
         self.mgr.add_files("session", "widget", [file2])
         self.assertEqual([file1, file2], self.mgr.get_files("session", "widget"))
         self.assertEqual(len(self.filemgr_events), 2)