Merge pull request #258 from ntabris/preserve-url-query

Append any URL query args to dashboard URL paths
dask · Feb 19, 2023 · f638295 · f638295
2 parents 4c23d29 + 55d0bb1
commit f638295
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/dask_labextension/dashboardhandler.py b/dask_labextension/dashboardhandler.py
@@ -29,14 +29,22 @@ async def get(self, url) -> None:
         try:
             client = httpclient.AsyncHTTPClient()
 
+            # Extract query (if any) from URL, this will then be appended after path.
+            # This allows using (eg) "?token=[...]" in URL for authentication.
+            if "?" in url:
+                pos = url.find("?")
+                url, query = url[:pos], url[pos:]
+            else:
+                query = ""
+
             # First check for the individual-plots endpoint at user-provided url.
             # We don't check for the root URL because that can trigger a lot of
             # object creation in the bokeh document.
             url = _normalize_dashboard_link(parse.unquote(url), self.request)
             effective_url = None
             individual_plots_url = url_path_join(
                 url,
-                "individual-plots.json",
+                f"individual-plots.json{query}",
             )
             try:
                 self.log.debug(
@@ -73,6 +81,14 @@ async def get(self, url) -> None:
 
             individual_plots = json.loads(individual_plots_response.body)
 
+            # If there was query in original URL, append to URLs returned
+            if query:
+                for name, plot_url in individual_plots.items():
+                    individual_plots[name] = f"{plot_url}{query}"
+                url = f"{url}{query}"
+                if effective_url:
+                    effective_url = f"{effective_url}{query}"
+
             self.set_status(200)
             self.finish(
                 json.dumps(