Skip to content

Commit

Permalink
Fix upload cancellation (#4527)
Browse files Browse the repository at this point in the history
* Improve assertions in test_cancel_upload

* Fix upload cancellation by using `refs` instead of `upload_controllers`
  • Loading branch information
masenf authored Dec 12, 2024
1 parent adfda8a commit 2d9849e
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 11 deletions.
11 changes: 5 additions & 6 deletions reflex/.templates/web/utils/state.js
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,6 @@ let event_processing = false;
// Array holding pending events to be processed.
const event_queue = [];

// Pending upload promises, by id
const upload_controllers = {};

/**
* Generate a UUID (Used for session tokens).
* Taken from: https://stackoverflow.com/questions/105034/how-do-i-create-a-guid-uuid
Expand Down Expand Up @@ -486,7 +483,9 @@ export const uploadFiles = async (
return false;
}

if (upload_controllers[upload_id]) {
const upload_ref_name = `__upload_controllers_${upload_id}`

if (refs[upload_ref_name]) {
console.log("Upload already in progress for ", upload_id);
return false;
}
Expand Down Expand Up @@ -546,7 +545,7 @@ export const uploadFiles = async (
});

// Send the file to the server.
upload_controllers[upload_id] = controller;
refs[upload_ref_name] = controller;

try {
return await axios.post(getBackendURL(UPLOADURL), formdata, config);
Expand All @@ -566,7 +565,7 @@ export const uploadFiles = async (
}
return false;
} finally {
delete upload_controllers[upload_id];
delete refs[upload_ref_name];
}
};

Expand Down
8 changes: 5 additions & 3 deletions reflex/components/core/upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from reflex.utils import format
from reflex.utils.imports import ImportVar
from reflex.vars import VarData
from reflex.vars.base import CallableVar, LiteralVar, Var, get_unique_variable_name
from reflex.vars.base import CallableVar, Var, get_unique_variable_name
from reflex.vars.sequence import LiteralStringVar

DEFAULT_UPLOAD_ID: str = "default"
Expand Down Expand Up @@ -108,7 +108,8 @@ def clear_selected_files(id_: str = DEFAULT_UPLOAD_ID) -> EventSpec:
# UploadFilesProvider assigns a special function to clear selected files
# into the shared global refs object to make it accessible outside a React
# component via `run_script` (otherwise backend could never clear files).
return run_script(f"refs['__clear_selected_files']({id_!r})")
func = Var("__clear_selected_files")._as_ref()
return run_script(f"{func}({id_!r})")


def cancel_upload(upload_id: str) -> EventSpec:
Expand All @@ -120,7 +121,8 @@ def cancel_upload(upload_id: str) -> EventSpec:
Returns:
An event spec that cancels the upload when triggered.
"""
return run_script(f"upload_controllers[{LiteralVar.create(upload_id)!s}]?.abort()")
controller = Var(f"__upload_controllers_{upload_id}")._as_ref()
return run_script(f"{controller}?.abort()")


def get_upload_dir() -> Path:
Expand Down
17 changes: 15 additions & 2 deletions tests/integration/test_upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -381,9 +381,22 @@ async def test_cancel_upload(tmp_path, upload_file: AppHarness, driver: WebDrive
await asyncio.sleep(0.3)
cancel_button.click()

# look up the backend state and assert on progress
# Wait a bit for the upload to get cancelled.
await asyncio.sleep(0.5)

# Get interim progress dicts saved in the on_upload_progress handler.
async def _progress_dicts():
state = await upload_file.get_state(substate_token)
return state.substates[state_name].progress_dicts

# We should have _some_ progress
assert await AppHarness._poll_for_async(_progress_dicts)

# But there should never be a final progress record for a cancelled upload.
for p in await _progress_dicts():
assert p["progress"] != 1

state = await upload_file.get_state(substate_token)
assert state.substates[state_name].progress_dicts
file_data = state.substates[state_name]._file_data
assert isinstance(file_data, dict)
normalized_file_data = {Path(k).name: v for k, v in file_data.items()}
Expand Down

0 comments on commit 2d9849e

Please sign in to comment.