From 3ee9e11ce1bfc2f74ef8bfcd2ee0cc6a3f2249ef Mon Sep 17 00:00:00 2001
From: Elizabeth Thompson <eschutho@gmail.com>
Date: Mon, 15 Nov 2021 12:47:40 -0800
Subject: [PATCH] feat: add a config to enable retina quality images in
 screenshots (#17409)

* add feature flag for retina screenshot support

* use config for pixel density

* run black
---
 superset/config.py          |  6 +++++-
 superset/utils/webdriver.py | 15 +++++++++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index 4b571dad27..05d7b8ecd7 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1085,7 +1085,11 @@ EMAIL_REPORTS_USER = "admin"
 WEBDRIVER_TYPE = "firefox"
 
 # Window size - this will impact the rendering of the data
-WEBDRIVER_WINDOW = {"dashboard": (1600, 2000), "slice": (3000, 1200)}
+WEBDRIVER_WINDOW = {
+    "dashboard": (1600, 2000),
+    "slice": (3000, 1200),
+    "pixel_density": 1,
+}
 
 # An optional override to the default auth hook used to provide auth to the
 # offline webdriver
diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
index 8ce7ee7c25..9cedc0f934 100644
--- a/superset/utils/webdriver.py
+++ b/superset/utils/webdriver.py
@@ -27,7 +27,7 @@ from selenium.common.exceptions import (
     TimeoutException,
     WebDriverException,
 )
-from selenium.webdriver import chrome, firefox
+from selenium.webdriver import chrome, firefox, FirefoxProfile
 from selenium.webdriver.common.by import By
 from selenium.webdriver.remote.webdriver import WebDriver
 from selenium.webdriver.support import expected_conditions as EC
@@ -51,22 +51,26 @@ class DashboardStandaloneMode(Enum):
 
 
 class WebDriverProxy:
-    def __init__(
-        self, driver_type: str, window: Optional[WindowSize] = None,
-    ):
+    def __init__(self, driver_type: str, window: Optional[WindowSize] = None):
         self._driver_type = driver_type
         self._window: WindowSize = window or (800, 600)
         self._screenshot_locate_wait = current_app.config["SCREENSHOT_LOCATE_WAIT"]
         self._screenshot_load_wait = current_app.config["SCREENSHOT_LOAD_WAIT"]
 
     def create(self) -> WebDriver:
+        pixel_density = current_app.config["WEBDRIVER_WINDOW"].get("pixel_density", 1)
         if self._driver_type == "firefox":
             driver_class = firefox.webdriver.WebDriver
             options = firefox.options.Options()
+            profile = FirefoxProfile()
+            profile.set_preference("layout.css.devPixelsPerPx", str(pixel_density))
+            kwargs: Dict[Any, Any] = dict(options=options, firefox_profile=profile)
         elif self._driver_type == "chrome":
             driver_class = chrome.webdriver.WebDriver
             options = chrome.options.Options()
+            options.add_argument(f"--force-device-scale-factor={pixel_density}")
             options.add_argument(f"--window-size={self._window[0]},{self._window[1]}")
+            kwargs = dict(options=options)
         else:
             raise Exception(f"Webdriver name ({self._driver_type}) not supported")
         # Prepare args for the webdriver init
@@ -75,7 +79,6 @@ class WebDriverProxy:
         for arg in current_app.config["WEBDRIVER_OPTION_ARGS"]:
             options.add_argument(arg)
 
-        kwargs: Dict[Any, Any] = dict(options=options)
         kwargs.update(current_app.config["WEBDRIVER_CONFIGURATION"])
         logger.info("Init selenium driver")
 
@@ -102,7 +105,7 @@ class WebDriverProxy:
             pass
 
     def get_screenshot(
-        self, url: str, element_name: str, user: "User",
+        self, url: str, element_name: str, user: "User"
     ) -> Optional[bytes]:
         params = {"standalone": DashboardStandaloneMode.REPORT.value}
         req = PreparedRequest()