feat: Allow extension-less file downloads. #1627 (#1635)

h2oai · Dec 9, 2022 · 39e1af1 · 39e1af1
1 parent 2b16927
commit 39e1af1
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 26 deletions.
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -39,19 +39,6 @@
         "bdist_wheel",
       ]
     },
-    {
-      "name": "Debug PY App",
-      "type": "python",
-      "request": "launch",
-      "program": "${workspaceFolder}/py/venv/bin/wave",
-      "python": "${workspaceFolder}/py/venv/bin/python",
-      "cwd": "${workspaceFolder}/py/examples",
-      "args": [
-        "run",
-        "tour",
-        "--no-reload",
-      ],
-    },
     {
       "name": "Debug Py Tests",
       "type": "python",
@@ -104,6 +91,19 @@
         "${workspaceFolder}/tools/wavegen/build/**/*.js"
       ]
     },
+    {
+      "name": "Debug Wave App",
+      "type": "python",
+      "request": "launch",
+      "program": "${workspaceFolder}/py/venv/bin/wave",
+      "python": "${workspaceFolder}/py/venv/bin/python",
+      "cwd": "${workspaceFolder}/py/examples",
+      "args": [
+        "run",
+        "tour",
+        "--no-reload",
+      ],
+    },
     {
       "name": "Debug Wave Server",
       "type": "go",

diff --git a/file_server.go b/file_server.go
@@ -69,14 +69,17 @@ func (fs *FileServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
-		if path.Ext(r.URL.Path) == "" { // ignore requests for directories and ext-less files
+		trimmedPrefix := strings.TrimPrefix(r.URL.Path, fs.baseURL)
+		fsDirPath := path.Join(fs.dir, trimmedPrefix)
+		// Ignore requests for directories and non-existent / unaccessible files.
+		if fileInfo, err := os.Stat(filepath.FromSlash(fsDirPath)); err != nil || fileInfo.IsDir() {
 			echo(Log{"t": "file_download", "path": r.URL.Path, "error": "not found"})
 			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
 			return
 		}
 
 		echo(Log{"t": "file_download", "path": r.URL.Path})
-		r.URL.Path = strings.TrimPrefix(r.URL.Path, fs.baseURL) // public
+		r.URL.Path = trimmedPrefix // public
 		fs.handler.ServeHTTP(w, r)
 
 	case http.MethodPost:

diff --git a/py/h2o_wave/core.py b/py/h2o_wave/core.py
@@ -719,9 +719,12 @@ def download(self, url: str, path: str) -> str:
         # If path is a directory, get basename from url
         filepath = os.path.join(path, os.path.basename(url)) if os.path.isdir(path) else path
 
-        with open(filepath, 'wb') as f:
-            with self._http.stream('GET', f'{_config.hub_host_address}{url}') as r:
-                for chunk in r.iter_bytes():
+        with self._http.stream('GET', f'{_config.hub_host_address}{url}') as res:
+            if res.status_code != 200:
+                res.read()
+                raise ServiceError(f'Download failed (code={res.status_code}): {res.text}')
+            with open(filepath, 'wb') as f:
+                for chunk in res.iter_bytes():
                     f.write(chunk)
 
         return filepath
@@ -893,10 +896,12 @@ async def download(self, url: str, path: str) -> str:
         path = os.path.abspath(path)
         # If path is a directory, get basename from url
         filepath = os.path.join(path, os.path.basename(url)) if os.path.isdir(path) else path
-
-        with open(filepath, 'wb') as f:
-            async with self._http.stream('GET', f'{_config.hub_host_address}{url}') as r:
-                async for chunk in r.aiter_bytes():
+        async with self._http.stream('GET', f'{_config.hub_host_address}{url}') as res:
+            if res.status_code != 200:
+                await res.aread()
+                raise ServiceError(f'Download failed (code={res.status_code}): {res.text}')
+            with open(filepath, 'wb') as f:
+                async for chunk in res.aiter_bytes():
                     f.write(chunk)
 
         return filepath

diff --git a/py/tests/test_python_server.py b/py/tests/test_python_server.py
@@ -22,6 +22,8 @@
 
 base_url = os.getenv('H2O_WAVE_BASE_URL', '/')
 
+
+# TODO: Add cleanup (site.unload) to tests that upload files.
 class TestPythonServer(unittest.TestCase):
     def test_new_empty_card(self):
         page = site['/test']
@@ -383,7 +385,7 @@ def test_multipart_server(self):
 
     def test_upload_dir(self):
         upload_path, = site.upload_dir(os.path.join('tests', 'test_folder'))
-        download_path = site.download(f'{base_url}{upload_path}test.txt', 'test.txt')
+        download_path = site.download(f'{upload_path}/test.txt', 'test.txt')
         txt = read_file(download_path)
         os.remove(download_path)
         assert len(txt) > 0

diff --git a/py/tests/test_python_server_async.py b/py/tests/test_python_server_async.py
@@ -19,7 +19,7 @@
 
 from .utils import read_file
 
-
+# TODO: Add cleanup (site.unload) to tests that upload files.
 class TestPythonServerAsync(unittest.IsolatedAsyncioTestCase):
     def __init__(self, methodName: str = ...) -> None:
         super().__init__(methodName)
@@ -68,8 +68,7 @@ async def test_multipart_server(self):
 
     async def test_upload_dir(self):
         upload_path, = await self.site.upload_dir(os.path.join('tests', 'test_folder'))
-        base_url = os.getenv('H2O_WAVE_BASE_URL', '/')
-        download_path = await self.site.download(f'{base_url}{upload_path}test.txt', 'test.txt')
+        download_path = await self.site.download(f'{upload_path}/test.txt', 'test.txt')
         txt = read_file(download_path)
         os.remove(download_path)
         assert len(txt) > 0