From f5a6b12e23fd188f1acd5bb00ed95ce133ec4d98 Mon Sep 17 00:00:00 2001 From: David Butenhof Date: Thu, 1 Jun 2023 13:41:54 -0400 Subject: [PATCH] Touch up Content-Disposition header (#3439) * Touch up Content-Disposition header PBENCH-1169 Help clients with the output of `GET /datasets//inventory/` by "tuning" the `Content-Disposition` header automatically generated by Werkzeug's `send_file` API. I started thinking I would construct the header and looking into how to add headers to the `Response` object constructed by `send_file`. It actually does a good job of guessing the filename, but I decided to help it by providing the name we want, and I'm designating the raw tarfile as an `attachment` that we can't reasonably expect a client (especially the UI) to do anything with "inline" (the default). --- .../api/resources/datasets_inventory.py | 15 +++++++- .../unit/server/test_datasets_inventory.py | 35 ++++++++++++------- 2 files changed, 37 insertions(+), 13 deletions(-) diff --git a/lib/pbench/server/api/resources/datasets_inventory.py b/lib/pbench/server/api/resources/datasets_inventory.py index fa46c887af..16aa72ab7d 100644 --- a/lib/pbench/server/api/resources/datasets_inventory.py +++ b/lib/pbench/server/api/resources/datasets_inventory.py @@ -75,7 +75,20 @@ def _get( file_path = dataset_location / target if file_path.is_file(): - return send_file(file_path) + # Tell send_file to set `Content-Disposition` to "attachment" if + # targeting the large binary tarball. Otherwise we'll recommend the + # default "inline": only `is_file()` paths are allowed here, and + # most are likely "displayable". While `send_file` will guess the + # download_name from the path, setting it explicitly does no harm + # and supports a unit test mock with no real file. + # + # NOTE: we could choose to be "smarter" based on file size, file + # type codes (e.g., .csv, .json), and so forth. + return send_file( + file_path, + as_attachment=target is None, + download_name=file_path.name, + ) elif file_path.exists(): raise APIAbort( HTTPStatus.UNSUPPORTED_MEDIA_TYPE, diff --git a/lib/pbench/test/unit/server/test_datasets_inventory.py b/lib/pbench/test/unit/server/test_datasets_inventory.py index 5248b9043a..c48319eb22 100644 --- a/lib/pbench/test/unit/server/test_datasets_inventory.py +++ b/lib/pbench/test/unit/server/test_datasets_inventory.py @@ -1,9 +1,9 @@ from http import HTTPStatus from pathlib import Path +from typing import Any, Optional import pytest import requests -import werkzeug.utils from pbench.server.cache_manager import CacheManager from pbench.server.database.models.datasets import Dataset, DatasetNotFound @@ -44,7 +44,7 @@ def query_api( def mock_find_dataset(self, dataset): class Tarball(object): unpacked = Path("/dataset/") - tarball_path = Path("/dataset_tarball") + tarball_path = Path("/dataset/tarball.tar.xz") # Validate the resource_id Dataset.query(resource_id=dataset) @@ -103,34 +103,45 @@ def test_not_a_file(self, query_get_as, monkeypatch): } def test_dataset_in_given_path(self, query_get_as, monkeypatch): - file_sent = None + mock_args: Optional[tuple[Path, dict[str, Any]]] = None def mock_send_file(path_or_file, *args, **kwargs): - nonlocal file_sent - file_sent = path_or_file + nonlocal mock_args + mock_args = (path_or_file, kwargs) return {"status": "OK"} monkeypatch.setattr(CacheManager, "find_dataset", self.mock_find_dataset) monkeypatch.setattr(Path, "is_file", lambda self: True) - monkeypatch.setattr(werkzeug.utils, "send_file", mock_send_file) + monkeypatch.setattr( + "pbench.server.api.resources.datasets_inventory.send_file", mock_send_file + ) response = query_get_as("fio_2", "1-default/default.csv", HTTPStatus.OK) assert response.status_code == HTTPStatus.OK - assert str(file_sent) == "/dataset/1-default/default.csv" + + path, args = mock_args + assert str(path) == "/dataset/1-default/default.csv" + assert args["as_attachment"] is False + assert args["download_name"] == "default.csv" @pytest.mark.parametrize("key", (None, "")) def test_get_result_tarball(self, query_get_as, monkeypatch, key): - file_sent = None + mock_args: Optional[tuple[Path, dict[str, Any]]] = None def mock_send_file(path_or_file, *args, **kwargs): - nonlocal file_sent - file_sent = path_or_file + nonlocal mock_args + mock_args = (path_or_file, kwargs) return {"status": "OK"} monkeypatch.setattr(CacheManager, "find_dataset", self.mock_find_dataset) monkeypatch.setattr(Path, "is_file", lambda self: True) - monkeypatch.setattr(werkzeug.utils, "send_file", mock_send_file) + monkeypatch.setattr( + "pbench.server.api.resources.datasets_inventory.send_file", mock_send_file + ) response = query_get_as("fio_2", key, HTTPStatus.OK) assert response.status_code == HTTPStatus.OK - assert str(file_sent) == "/dataset_tarball" + path, args = mock_args + assert str(path) == "/dataset/tarball.tar.xz" + assert args["as_attachment"] is True + assert args["download_name"] == "tarball.tar.xz"