From bda43ac0f6be3523230ccd5703b57547e8aed6f3 Mon Sep 17 00:00:00 2001
From: Kamil Gabryjelski <kamil.gabryjelski@gmail.com>
Date: Fri, 27 Oct 2023 10:37:32 +0200
Subject: [PATCH] chore: Add config options for Playwright wait_until and
 default timeout (#25765)

---
 superset/config.py          |  8 ++++++++
 superset/utils/webdriver.py | 24 ++++++++++++------------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index 936c8d6d4b6fb..dd244dc14a42a 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -682,6 +682,14 @@ class D3Format(TypedDict, total=False):
 SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE = 5
 # Max time to wait for error message modal to close, in seconds
 SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE = 5
+# Event that Playwright waits for when loading a new page
+# Possible values: "load", "commit", "domcontentloaded", "networkidle"
+# Docs: https://playwright.dev/python/docs/api/class-page#page-goto-option-wait-until
+SCREENSHOT_PLAYWRIGHT_WAIT_EVENT = "load"
+# Default timeout for Playwright browser context for all operations
+SCREENSHOT_PLAYWRIGHT_DEFAULT_TIMEOUT = int(
+    timedelta(seconds=30).total_seconds() * 1000
+)
 
 # ---------------------------------------------------
 # Image and file configuration
diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
index 720c399b2a67a..4353319072287 100644
--- a/superset/utils/webdriver.py
+++ b/superset/utils/webdriver.py
@@ -49,7 +49,7 @@
     from playwright.sync_api import (
         BrowserContext,
         ElementHandle,
-        Error,
+        Error as PlaywrightError,
         Page,
         sync_playwright,
         TimeoutError as PlaywrightTimeout,
@@ -140,9 +140,9 @@ def find_unexpected_errors(page: Page) -> list[str]:
                         "(node, error_html) => node.innerHtml = error_html",
                         [error_as_html],
                     )
-                except Error:
+                except PlaywrightError:
                     logger.exception("Failed to update error messages using alert_div")
-        except Error:
+        except PlaywrightError:
             logger.exception("Failed to capture unexpected errors")
 
         return error_messages
@@ -161,9 +161,14 @@ def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | Non
                 },
                 device_scale_factor=pixel_density,
             )
+            context.set_default_timeout(
+                current_app.config["SCREENSHOT_PLAYWRIGHT_DEFAULT_TIMEOUT"]
+            )
             self.auth(user, context)
             page = context.new_page()
-            page.goto(url)
+            page.goto(
+                url, wait_until=current_app.config["SCREENSHOT_PLAYWRIGHT_WAIT_EVENT"]
+            )
             img: bytes | None = None
             selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"]
             logger.debug("Sleeping for %i seconds", selenium_headstart)
@@ -202,7 +207,7 @@ def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | Non
                     )
                     page.wait_for_selector(
                         ".loading",
-                        timeout=self._screenshot_locate_wait * 1000,
+                        timeout=self._screenshot_load_wait * 1000,
                         state="detached",
                     )
                 except PlaywrightTimeout as ex:
@@ -236,14 +241,9 @@ def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | Non
             except PlaywrightTimeout:
                 # raise again for the finally block, but handled above
                 pass
-            except StaleElementReferenceException:
-                logger.exception(
-                    "Selenium got a stale element while requesting url %s",
-                    url,
-                )
-            except WebDriverException:
+            except PlaywrightError:
                 logger.exception(
-                    "Encountered an unexpected error when requeating url %s", url
+                    "Encountered an unexpected error when requesting url %s", url
                 )
             return img