Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 28 additions & 9 deletions superset/utils/screenshot_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,12 +207,16 @@ def resolve_screenshot_task_budget_seconds(
"""


def combine_screenshot_tiles(screenshot_tiles: list[bytes]) -> bytes:
def combine_screenshot_tiles(
screenshot_tiles: list[bytes], log_context: str | None = None
) -> bytes:
"""
Combine multiple screenshot tiles into a single vertical image.

Args:
screenshot_tiles: List of screenshot bytes in PNG format
log_context: Optional identifier (e.g. report execution id, or a
thumbnail cache key) appended to log lines for tracing.

Returns:
Combined screenshot as bytes
Expand All @@ -223,6 +227,7 @@ def combine_screenshot_tiles(screenshot_tiles: list[bytes]) -> bytes:
if len(screenshot_tiles) == 1:
return screenshot_tiles[0]

context_suffix = f" [{log_context}]" if log_context else ""
try:
# Open all images
images = [Image.open(io.BytesIO(tile)) for tile in screenshot_tiles]
Expand All @@ -246,7 +251,7 @@ def combine_screenshot_tiles(screenshot_tiles: list[bytes]) -> bytes:
return output.getvalue()

except Exception as e:
logger.exception("Failed to combine screenshot tiles: %s", e)
logger.exception("Failed to combine screenshot tiles: %s%s", e, context_suffix)
# Return the first tile as fallback
return screenshot_tiles[0]

Expand Down Expand Up @@ -309,16 +314,17 @@ def take_tiled_screenshot(
dashboard_top = element_info["top"]

logger.info(
"Dashboard: %sx%spx at (%s, %s)",
"Dashboard: %sx%spx at (%s, %s)%s",
dashboard_width,
dashboard_height,
dashboard_left,
dashboard_top,
context_suffix,
)

# Calculate number of tiles needed
num_tiles = max(1, (dashboard_height + tile_height - 1) // tile_height)
logger.info("Taking %s screenshot tiles", num_tiles)
logger.info("Taking %s screenshot tiles%s", num_tiles, context_suffix)

screenshot_tiles = []

Expand All @@ -328,7 +334,11 @@ def take_tiled_screenshot(

page.evaluate(f"window.scrollTo(0, {scroll_y})")
logger.debug(
"Scrolled window to %s for tile %s/%s", scroll_y, i + 1, num_tiles
"Scrolled window to %s for tile %s/%s%s",
scroll_y,
i + 1,
num_tiles,
context_suffix,
)
# Wait for scroll to settle and content to load
page.wait_for_timeout(SCROLL_SETTLE_TIMEOUT_MS)
Expand Down Expand Up @@ -401,13 +411,14 @@ def take_tiled_screenshot(
logger.warning(
"Skipping tile %s/%s due to invalid clip dimensions: "
"x=%s, y=%s, width=%s, height=%s "
"(element may be scrolled out of viewport)",
"(element may be scrolled out of viewport).%s",
i + 1,
num_tiles,
clip_x,
clip_y,
dashboard_width,
clip_height,
context_suffix,
)
continue

Expand All @@ -423,11 +434,19 @@ def take_tiled_screenshot(
tile_screenshot = page.screenshot(type="png", clip=clip)
screenshot_tiles.append(tile_screenshot)

logger.debug("Captured tile %s/%s with clip %s", i + 1, num_tiles, clip)
logger.debug(
"Captured tile %s/%s with clip %s%s",
i + 1,
num_tiles,
clip,
context_suffix,
)

# Combine all tiles
logger.info("Combining screenshot tiles...")
combined_screenshot = combine_screenshot_tiles(screenshot_tiles)
logger.info("Combining screenshot tiles...%s", context_suffix)
combined_screenshot = combine_screenshot_tiles(
screenshot_tiles, log_context=log_context
)

return combined_screenshot

Expand Down
54 changes: 42 additions & 12 deletions superset/utils/screenshots.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,10 @@ def __init__(self, url: str, digest: str | None):
self.screenshot = None

def driver(
self, window_size: WindowSize | None = None, user: User | None = None
self,
window_size: WindowSize | None = None,
user: User | None = None,
log_context: str | None = None,
) -> WebDriverProxy:
window_size = window_size or self.window_size
if feature_flag_manager.is_feature_enabled("PLAYWRIGHT_REPORTS_AND_THUMBNAILS"):
Expand All @@ -199,11 +202,13 @@ def driver(
return WebDriverPlaywright(self.driver_type, window_size)

# Playwright not available, falling back to Selenium
context_suffix = f" [{log_context}]" if log_context else ""
logger.info(
"PLAYWRIGHT_REPORTS_AND_THUMBNAILS enabled but Playwright not "
"installed. Falling back to Selenium (WebGL/Canvas charts may "
"not render correctly). %s",
"not render correctly). %s%s",
PLAYWRIGHT_INSTALL_MESSAGE,
context_suffix,
)

# Use Selenium as default/fallback
Expand All @@ -215,7 +220,7 @@ def get_screenshot(
window_size: WindowSize | None = None,
log_context: str | None = None,
) -> bytes | None:
driver = self.driver(window_size, user)
driver = self.driver(window_size, user, log_context=log_context)
try:
self.screenshot = driver.get_screenshot(
self.url, self.element, user, log_context=log_context
Expand Down Expand Up @@ -310,22 +315,38 @@ def compute_and_cache( # pylint: disable=too-many-arguments
image = None
# Assuming all sorts of things can go wrong with Selenium
try:
logger.info("trying to generate screenshot")
logger.info(
"trying to generate screenshot for cache_key=%s", cache_key
)
with event_logger.log_context(
f"screenshot.compute.{self.thumbnail_type}"
):
image = self.get_screenshot(user=user, window_size=window_size)
image = self.get_screenshot(
user=user,
window_size=window_size,
log_context=f"cache_key={cache_key}",
)
except Exception as ex: # pylint: disable=broad-except
logger.warning(
"Failed at generating thumbnail %s", ex, exc_info=True
"Failed at generating thumbnail for cache_key=%s: %s",
cache_key,
ex,
exc_info=True,
)
cache_payload.error()
if image and window_size != thumb_size:
try:
image = self.resize_image(image, thumb_size=thumb_size)
image = self.resize_image(
image,
thumb_size=thumb_size,
log_context=f"cache_key={cache_key}",
)
except Exception as ex: # pylint: disable=broad-except
logger.warning(
"Failed at resizing thumbnail %s", ex, exc_info=True
"Failed at resizing thumbnail for cache_key=%s: %s",
cache_key,
ex,
exc_info=True,
)
cache_payload.error()
image = None
Expand All @@ -344,7 +365,9 @@ def compute_and_cache( # pylint: disable=too-many-arguments
logger.info("Caching thumbnail: %s", cache_key)
self.cache.set(cache_key, cache_payload.to_dict())
logger.info(
"Updated thumbnail cache; Status: %s", cache_payload.get_status()
"Updated thumbnail cache for %s; Status: %s",
cache_key,
cache_payload.get_status(),
)
except LockAlreadyHeldException:
logger.info(
Expand All @@ -359,16 +382,23 @@ def resize_image(
output: str = "png",
thumb_size: WindowSize | None = None,
crop: bool = True,
log_context: str | None = None,
) -> bytes:
context_suffix = f" [{log_context}]" if log_context else ""
thumb_size = thumb_size or cls.thumb_size
img = Image.open(BytesIO(img_bytes))
logger.debug("Selenium image size: %s", str(img.size))
logger.debug("Selenium image size: %s%s", str(img.size), context_suffix)
if crop and img.size[1] != cls.window_size[1]:
desired_ratio = float(cls.window_size[1]) / cls.window_size[0]
desired_width = int(img.size[0] * desired_ratio)
logger.debug("Cropping to: %s*%s", str(img.size[0]), str(desired_width))
logger.debug(
"Cropping to: %s*%s%s",
str(img.size[0]),
str(desired_width),
context_suffix,
)
img = img.crop((0, 0, img.size[0], desired_width))
logger.debug("Resizing to %s", str(thumb_size))
logger.debug("Resizing to %s%s", str(thumb_size), context_suffix)
img = img.resize(thumb_size, Image.Resampling.LANCZOS)
new_img = BytesIO()
if output != "png":
Expand Down
Loading
Loading