From c9fa363ba4d4adfdfe2df4a63fa92ca65b1206d1 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 29 Jul 2026 15:42:33 +0000 Subject: [PATCH] fix: capture S3 diagnostics and stop logging presigned URLs on SQL cache failures `upload_sql_cache` reported failures using only the string from `raise_for_status()`. For presigned S3 uploads that discarded the one thing identifying the cause (the response body) while including the one thing that should never reach logs (the presigned URL, which carries X-Amz-Credential and X-Amz-Security-Token). The unique URL in every message also meant no two occurrences shared a message string, so downstream error tracking could not group them. - Extract S3's /// plus x-amz-request-id rather than the exception string. Only those four elements are read; , and echo the signed query string and are never logged. - Redact credential-bearing material from any text destined for a log. The primary defence strips the query string off anything URL-shaped, including the scheme-less path-only form urllib3 puts in connection errors - so network failures, not just HTTP ones, are covered. - Carry the URL's issue time in a SqlCacheUpload NamedTuple and log seconds_since_url_issued beside url_expires_in, making an expiry-driven failure self-evident. - Use constant log messages with all variable data in `extra`, at all four sites in the module. - Give serialization failures a distinct `failed_to_serialize_cache` cause, and log only the exception type for them - the message quotes user column names. - Fetch the cached object explicitly instead of handing the URL to pandas, which fetched it with urllib and raised before S3's error body was ever read. The parquet/pickle fallback is preserved and now downloads once instead of twice. - Bound the connect phase of both transfers at 5s; read is left unbounded so slow but healthy transfers of large cache objects are unaffected. Cache failures remain swallowed and can never fail the user's query. tests/unit/test_sql_caching.py goes from 26 tests in 31.1s to 74 in 1.4s - the former suite had a test making a real network call to localhost:19456 whose mocks were provably dead. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_015xJczfWE4mq3ux7ZU9ZeU9 --- deepnote_toolkit/sql/sql_caching.py | 401 ++++++++++-- deepnote_toolkit/sql/sql_execution.py | 12 +- tests/unit/test_sql_caching.py | 848 +++++++++++++++++++++++++- 3 files changed, 1183 insertions(+), 78 deletions(-) diff --git a/deepnote_toolkit/sql/sql_caching.py b/deepnote_toolkit/sql/sql_caching.py index db396e9..3f325b6 100644 --- a/deepnote_toolkit/sql/sql_caching.py +++ b/deepnote_toolkit/sql/sql_caching.py @@ -1,6 +1,11 @@ import hashlib import json +import re import tempfile +import time +from io import BytesIO +from typing import IO, Any, NamedTuple, Optional +from urllib.parse import parse_qs, urlsplit import pandas as pd import requests @@ -15,10 +20,239 @@ # Initialize logger logger = get_logger() +# Only the connect phase is bounded. A read timeout would abort slow but healthy +# transfers of large cache objects, which would be a regression. +_OBJECT_STORE_TIMEOUT: tuple[int, None] = (5, None) + +# Bounds on every piece of third-party text that can reach a log entry +_MAX_ERROR_BODY_BYTES = 4096 +_MAX_ERROR_FIELD_CHARS = 500 +_MAX_OBJECT_PATH_CHARS = 200 +# An exception message is the one input that is not already bounded by the time it +# is redacted, so it is cut first. The margin over _MAX_ERROR_FIELD_CHARS leaves +# room for the text to grow as placeholders replace what they redact. +_MAX_RAW_EXCEPTION_CHARS = _MAX_ERROR_FIELD_CHARS * 4 + +# Strips the query string off anything URL-shaped, including the scheme-less, +# path-only form that urllib3 puts in its connection error messages. The character +# classes exclude whitespace and angle brackets, so a match can never span an XML +# tag boundary or swallow ordinary prose containing a question mark. +_URL_QUERY_PATTERN = re.compile( + r"((?:https?://)?[^\s?\"'<>]*/[^\s?\"'<>]*)\?[^\s\"'<>]*" +) +# Backstop for signing parameters that appear without a URL in front of them, as +# they do in the canonical request echoed by an S3 SignatureDoesNotMatch body +_AWS_CREDENTIAL_PARAM_PATTERN = re.compile( + r"(X-(?:Amz|Goog)-(?:Credential|Security-Token|Signature)=)[^&\s\"'<>]*", + re.IGNORECASE, +) +# S3 states the real cause of a failure in an XML body. and +# accompany a rejection for an elapsed signature and are AWS's own answer to +# whether the URL ran out of time, independent of this pod's clock. +_S3_ERROR_FIELD_PATTERN = re.compile( + r"<(Code|Message|Expires|ServerTime)>(.*?)", re.DOTALL | re.IGNORECASE +) + + +class SqlCacheUpload(NamedTuple): + """A presigned upload URL together with the moment it was issued. + + The URL is obtained before the query runs and used only after the query + completes and its result is serialized, so all of that has to fit inside the + URL's signed lifetime. Carrying issued_at is what makes that measurable. + """ + + url: str + # time.monotonic() taken just before the URL was requested from the webapp + issued_at: float + + +class _SqlCacheHttpError(Exception): + """A non-2xx response from the cache object store. + + Carries pre-computed, non-sensitive diagnostics for the caller to log. Its own + string representation deliberately contains no URL. + """ + + def __init__(self, diagnostics: dict[str, Any]) -> None: + """Store diagnostics already extracted from the failed response. + + Args: + diagnostics (dict): Log-safe fields describing the response. + """ + super().__init__("SQL cache object store returned an error response") + self.diagnostics = diagnostics + + +def _redact_sensitive(text: str) -> str: + """Strip credential-bearing material from text destined for a log. + + Dropping the query string from anything URL-shaped is the primary defence: + presigned URLs carry X-Amz-Credential and X-Amz-Security-Token there, and + urllib3 embeds the requested path and its query in every connection error. + Blanking named signing parameters is a backstop for the same values appearing + without a URL in front of them. + """ + redacted = _URL_QUERY_PATTERN.sub(r"\1?", text) + return _AWS_CREDENTIAL_PARAM_PATTERN.sub(r"\1", redacted) + + +def _to_int_or_none(value: Optional[str]) -> Optional[int]: + """Coerce a query string value to an int, or None when it is not numeric.""" + if value is None: + return None + + try: + return int(value) + except ValueError: + return None + + +def _seconds_since(issued_at: float) -> Optional[float]: + """Monotonic seconds elapsed since issued_at, or None if it is not a number. + + Guarded because it is evaluated while building a log entry, in a function + documented as never failing the user's query. + """ + try: + return round(time.monotonic() - issued_at, 1) + except Exception: + return None + + +def _safe_url_path(url: str) -> Optional[str]: + """The bounded path of a URL, or None when the URL cannot be parsed. + + The query string is never returned, and neither is the input itself when + urlsplit rejects it. + """ + if not isinstance(url, str): + # urlsplit(None) answers with bytes-valued fields instead of raising, and a + # bytes value in extra silently discards the whole error report + return None + + try: + return urlsplit(url).path[:_MAX_OBJECT_PATH_CHARS] + except Exception: + return None + + +def _describe_presigned_url(url: str) -> dict[str, Any]: + """Object path and declared expiry, without the credential-bearing query string. + + The query string of a presigned URL carries X-Amz-Credential and + X-Amz-Security-Token and must never reach a log, so it is never returned - not + even when the URL cannot be parsed. + """ + if not isinstance(url, str): + # urlsplit(None) answers with bytes-valued fields instead of raising, and a + # bytes value in extra silently discards the whole error report + return {"object_host": None, "object_path": None, "url_expires_in": None} + + try: + parts = urlsplit(url) + expires_values = parse_qs(parts.query).get("X-Amz-Expires") + return { + # hostname rather than netloc, which would carry any user:pass@ userinfo + "object_host": parts.hostname, + "object_path": parts.path[:_MAX_OBJECT_PATH_CHARS], + "url_expires_in": _to_int_or_none( + expires_values[0] if expires_values else None + ), + } + except Exception: + # Built purely from literals, so an unparseable URL cannot leak through + # the failure path either + return {"object_host": None, "object_path": None, "url_expires_in": None} + + +def _read_response_body_prefix(response: requests.Response) -> Optional[str]: + """Decode a bounded prefix of a response body, or None when it is unreadable. + + The prefix is sliced off the raw bytes before decoding, so an oversized or + mis-declared body costs nothing to inspect and binary content cannot raise. + """ + try: + return response.content[:_MAX_ERROR_BODY_BYTES].decode( + "utf-8", errors="replace" + ) + except Exception: + # A truncated or badly encoded transfer. The caller still has the status + # code, which is worth logging on its own. + return None + + +def _describe_s3_error(response: requests.Response) -> dict[str, Any]: + """Non-sensitive diagnostics from a failed response from the object store. + + Only , , and are taken from the body, + because the other elements of an S3 error document (, + , ) echo the signed query string and with it the + credentials this module must not log. + """ + diagnostics: dict[str, Any] = { + "status_code": response.status_code, + # Needed if we ever have to ask AWS Support about a specific request + "aws_request_id": response.headers.get("x-amz-request-id"), + "aws_host_id": response.headers.get("x-amz-id-2"), + # AWS's clock at the moment it rejected us + "aws_date": response.headers.get("Date"), + "s3_error_code": None, + "s3_error_message": None, + "s3_expires": None, + "s3_server_time": None, + } + + body = _read_response_body_prefix(response) + if body is None: + return diagnostics + + fields = { + name.lower(): value for name, value in _S3_ERROR_FIELD_PATTERN.findall(body) + } + for key, name in ( + ("s3_error_code", "code"), + ("s3_error_message", "message"), + ("s3_expires", "expires"), + ("s3_server_time", "servertime"), + ): + value = fields.get(name) + if value is not None: + diagnostics[key] = _redact_sensitive(value)[:_MAX_ERROR_FIELD_CHARS] + + if diagnostics["s3_error_code"] is None: + # Not an S3 error document - a proxy or gateway answered instead. Keep a + # short redacted snippet so that case stays diagnosable. + diagnostics["response_body_snippet"] = _redact_sensitive(body)[ + :_MAX_ERROR_FIELD_CHARS + ] + + return diagnostics + + +def _describe_exception(exc: BaseException) -> dict[str, Any]: + """Non-sensitive diagnostics from an exception raised while using the cache.""" + if isinstance(exc, _SqlCacheHttpError): + return dict(exc.diagnostics) + + # Cut before redacting rather than after: the URL pattern backtracks over every + # start position in a long run of non-separator characters, and an exception + # message has no bound of its own. Cutting a query string short is safe, since + # the pattern consumes to the end of the string either way. + message = _redact_sensitive(str(exc)[:_MAX_RAW_EXCEPTION_CHARS]) + return { + "error_type": type(exc).__name__, + "error_message": message[:_MAX_ERROR_FIELD_CHARS], + } + def get_sql_cache( - query, bind_params, integration_id, sql_cache_mode, return_variable_type -): + query: str, + bind_params: dict, + integration_id: str, + sql_cache_mode: str, + return_variable_type: str, +) -> tuple[Optional[pd.DataFrame], Optional[SqlCacheUpload]]: """ Retrieves the SQL cache from webapp for a given query. @@ -27,9 +261,11 @@ def get_sql_cache( bind_params (dict): The bind parameters for the SQL query. integration_id (str): The integration ID associated with the cache. sql_cache_mode (str): The mode of the SQL cache. + return_variable_type (str): The type of variable the result is bound to. Returns: - tuple: A tuple containing the cached dataframe (if available) and the upload URL (if applicable). + tuple: A tuple containing the cached dataframe (if available) and the + pending upload (if applicable). """ if not is_single_select_query(query): @@ -44,6 +280,11 @@ def get_sql_cache( query_hash = _generate_cache_key(query, bind_params) + # Taken before the request because the webapp signs the upload URL somewhere + # inside this round trip, which makes it a conservative upper bound on how much + # of the URL's lifetime has been spent by the time we come to use it + requested_at = time.monotonic() + cache_info = None try: cache_info = _request_cache_info_from_webapp( @@ -52,9 +293,11 @@ def get_sql_cache( except Exception as exc: # we failed to request the cache info from the webapp logger.error( - "Failed to request SQL cache info: %s", - exc, - extra={"sql_caching_cause": "failed_to_request_cache_info"}, + "Failed to request SQL cache info", + extra={ + "sql_caching_cause": "failed_to_request_cache_info", + **_describe_exception(exc), + }, ) return None, None @@ -67,9 +310,12 @@ def get_sql_cache( except Exception as exc: # we failed to download the dataframe from the cache logger.error( - "Failed to download dataframe from cache: %s", - exc, - extra={"sql_caching_cause": "failed_to_download_from_cache"}, + "Failed to download dataframe from cache", + extra={ + "sql_caching_cause": "failed_to_download_from_cache", + **_describe_exception(exc), + **_describe_presigned_url(download_url), + }, ) return None, None @@ -86,42 +332,98 @@ def get_sql_cache( return dataframe_from_cache, None if cache_info["result"] == "cacheMiss" or cache_info["result"] == "alwaysWrite": - return None, cache_info["uploadUrl"] + return None, SqlCacheUpload( + url=cache_info["uploadUrl"], issued_at=requested_at + ) return None, None -def upload_sql_cache(dataframe, upload_url): - """upload the result to the cache as a parquet file""" +def _serialize_dataframe_for_cache( + dataframe: pd.DataFrame, file_obj: IO[bytes] +) -> None: + """Write the dataframe to file_obj as parquet, falling back to pickle.""" + try: + dataframe.to_parquet(file_obj) + except (ArrowNotImplementedError, ArrowInvalid, OverflowError): + # see NB-1684 + # we fallback to pickle if parquet serialization fails (which will throw either of first 2 errors) + # OverflowError: PyArrow raises this for Python int / Decimal values exceeding int64 range + file_obj.seek(0) + file_obj.truncate() + dataframe.to_pickle(file_obj) + + +def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: + """Upload the result to the cache as a parquet file. + + Caching is best effort: every failure is logged under a constant message, with + all variable data in extra so occurrences group, and then swallowed so that it + can never fail the user's query. + + Args: + dataframe (pd.DataFrame): The query result to cache. + upload (SqlCacheUpload): The presigned upload URL and when it was issued. + """ try: with tempfile.TemporaryFile() as temp_file: try: - dataframe.to_parquet(temp_file) - except (ArrowNotImplementedError, ArrowInvalid, OverflowError): - # see NB-1684 - # we fallback to pickle if parquet serialization fails (which will throw either of first 2 errors) - # OverflowError: PyArrow raises this for Python int / Decimal values exceeding int64 range - temp_file.seek(0) - temp_file.truncate() - dataframe.to_pickle(temp_file) + _serialize_dataframe_for_cache(dataframe, temp_file) + except Exception as exc: + # Nothing was sent, so this is not an upload failure and gets its own + # cause. Only the type is logged: serialization errors quote the + # user's column names. + logger.error( + "Failed to upload SQL cache", + extra={ + "sql_caching_cause": "failed_to_serialize_cache", + "error_type": type(exc).__name__, + }, + ) + return temp_file.seek(0) - # PUT the file to cache_upload_url pre-signed s3 url - response = requests.put(upload_url, data=temp_file) - response.raise_for_status() + # PUT the file to the pre-signed s3 url + response = requests.put( + upload.url, data=temp_file, timeout=_OBJECT_STORE_TIMEOUT + ) + + if response.status_code < 400: + return + + # Checked explicitly rather than with raise_for_status(), whose message + # interpolates the whole presigned URL and discards S3's error body + failure = _describe_s3_error(response) except Exception as exc: - logger.error( - "Failed to upload SQL cache: %s", - exc, - extra={"sql_caching_cause": "failed_to_upload_to_cache"}, - ) + # The request never completed, so there is no response to describe + failure = _describe_exception(exc) + + logger.error( + "Failed to upload SQL cache", + extra={ + "sql_caching_cause": "failed_to_upload_to_cache", + **failure, + **_describe_presigned_url(upload.url), + "seconds_since_url_issued": _seconds_since(upload.issued_at), + }, + ) -def _try_read_cache(download_url): +def _try_read_cache(download_url: str) -> pd.DataFrame: + """Download the cached object and read it as a dataframe. + + The object is fetched explicitly instead of handing the URL to pandas, which + fetches it with urllib and so raises before S3's error body is ever read. + """ + response = requests.get(download_url, timeout=_OBJECT_STORE_TIMEOUT) + if response.status_code >= 400: + raise _SqlCacheHttpError(_describe_s3_error(response)) + + buffer = BytesIO(response.content) try: # Attempt to read as a parquet file - return pd.read_parquet(download_url) + return pd.read_parquet(buffer) except ArrowInvalid: # ArrowInvalid means that the file at download_url is not a parquet file. # We fallback to the pickle format if that happens, because the cache should either be in parquet or @@ -129,12 +431,10 @@ def _try_read_cache(download_url): # (see .to_pickle fallback in upload_sql_cache) pass - try: - # Attempt to read as a pickle file - return pd.read_pickle(download_url) - except Exception: - # If reading as pickle also fails, re-raise this exception to be caught by the caller - raise + # The failed parquet read left the cursor part-way through the buffer. Reading + # from the URL used to start a fresh download for each attempt. + buffer.seek(0) + return pd.read_pickle(buffer) def _generate_cache_key(query, bind_params): @@ -143,7 +443,19 @@ def _generate_cache_key(query, bind_params): ).hexdigest() -def _request_cache_info_from_webapp(query_hash, integration_id, sql_cache_mode): +def _request_cache_info_from_webapp( + query_hash: str, integration_id: str, sql_cache_mode: str +) -> Optional[dict[str, Any]]: + """Ask the webapp what the cache holds for this query. + + Args: + query_hash (str): The cache key derived from the query and its params. + integration_id (str): The integration ID associated with the cache. + sql_cache_mode (str): The mode of the SQL cache. + + Returns: + The cache info, or None when caching is disabled or unavailable. + """ # calls https://github.com/deepnote/deepnote/blob/eb96467937de12db8b588e5aa0a80244cec7eae7/apps/webapp/server/api/userpod-api.ts#L133 sql_cache_url = get_absolute_userpod_api_url( f"integrations/{integration_id}/sql-cache?sqlCacheKey={query_hash}&sqlCacheMode={sql_cache_mode}" @@ -158,8 +470,19 @@ def _request_cache_info_from_webapp(query_hash, integration_id, sql_cache_mode): ) if sql_cache_response.status_code != 200: # the caching endpoint is not available, we can't use it. We'll skip the caching logic - error_msg = f"Failed to request cache info from {sql_cache_url}, status code {sql_cache_response.status_code}, response {sql_cache_response.text}" - logger.error(error_msg, extra={"sql_caching_cause": "http_error"}) + body = _read_response_body_prefix(sql_cache_response) + snippet = ( + None if body is None else _redact_sensitive(body)[:_MAX_ERROR_FIELD_CHARS] + ) + logger.error( + "Failed to request cache info", + extra={ + "sql_caching_cause": "http_error", + "status_code": sql_cache_response.status_code, + "cache_info_path": _safe_url_path(sql_cache_url), + "response_body_snippet": snippet, + }, + ) return None result_dict = sql_cache_response.json() diff --git a/deepnote_toolkit/sql/sql_execution.py b/deepnote_toolkit/sql/sql_execution.py index 642a687..6f8daf8 100644 --- a/deepnote_toolkit/sql/sql_execution.py +++ b/deepnote_toolkit/sql/sql_execution.py @@ -439,10 +439,10 @@ def _execute_sql_with_caching( integration_id = sql_alchemy_dict.get("integration_id") can_get_sql_cache = integration_id is not None and sql_caching_enabled - cache_upload_url = None + cache_upload = None if can_get_sql_cache: - dataframe_from_cache, cache_upload_url = get_sql_cache( + dataframe_from_cache, cache_upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) if dataframe_from_cache is not None: @@ -461,7 +461,7 @@ def _execute_sql_with_caching( query_with_audit_comment, bind_params, sql_alchemy_dict, - cache_upload_url, + cache_upload, return_variable_type, query_preview_source, # The original query before any transformations such as appending a LIMIT clause ) @@ -492,7 +492,7 @@ def _query_data_source( query, bind_params, sql_alchemy_dict, - cache_upload_url, + cache_upload, return_variable_type, query_preview_source, ): @@ -536,8 +536,8 @@ def _query_data_source( # if df is larger than 5GB, don't upload it. See NB-988 dataframe_is_cacheable = dataframe_size_in_bytes < 5 * 1024 * 1024 * 1024 - if cache_upload_url is not None and dataframe_is_cacheable: - upload_sql_cache(dataframe, cache_upload_url) + if cache_upload is not None and dataframe_is_cacheable: + upload_sql_cache(dataframe, cache_upload) return dataframe finally: diff --git a/tests/unit/test_sql_caching.py b/tests/unit/test_sql_caching.py index 7e6f287..031fbfe 100644 --- a/tests/unit/test_sql_caching.py +++ b/tests/unit/test_sql_caching.py @@ -1,18 +1,219 @@ +import json +import logging +import time import unittest +from io import BytesIO from unittest import mock from unittest.mock import patch import pandas as pd +import requests from parameterized import parameterized from pyarrow import ArrowInvalid from deepnote_toolkit.sql.sql_caching import ( + _URL_QUERY_PATTERN, + SqlCacheUpload, + _describe_exception, + _describe_presigned_url, + _describe_s3_error, _generate_cache_key, + _redact_sensitive, + _request_cache_info_from_webapp, + _safe_url_path, get_sql_cache, upload_sql_cache, ) from deepnote_toolkit.sql.sql_utils import is_single_select_query +QUERY = "SELECT * FROM users" + +# A presigned URL of the shape the webapp hands us, with the credential-bearing +# parameters given values that are easy to search log output for +PRESIGNED_URL = ( + "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" + "?X-Amz-Algorithm=AWS4-HMAC-SHA256" + "&X-Amz-Credential=CREDVALUE%2F20260729%2Feu-west-1%2Fs3%2Faws4_request" + "&X-Amz-Date=20260729T120000Z&X-Amz-Expires=900" + "&X-Amz-Security-Token=TOKENVALUE&X-Amz-Signature=SIGVALUE" +) + +# What requests raises when it cannot reach the object store. The URL urllib3 +# embeds is path-only - there is no scheme and no host in front of it. +URLLIB3_ERROR_MESSAGE = ( + "HTTPSConnectionPool(host='bucket.s3.eu-west-1.amazonaws.com', port=443): " + "Max retries exceeded with url: /ws/int/key" + "?X-Amz-Algorithm=AWS4-HMAC-SHA256" + "&X-Amz-Credential=CREDVALUE%2F20260729%2Feu-west-1%2Fs3%2Faws4_request" + "&X-Amz-Date=20260729T120000Z&X-Amz-Expires=900" + "&X-Amz-Security-Token=TOKENVALUE&X-Amz-Signature=SIGVALUE " + "(Caused by NameResolutionError('Failed to resolve host'))" +) + +# Substrings that must never appear anywhere in a log entry from this module +SECRETS = ("CREDVALUE", "TOKENVALUE", "SIGVALUE", "X-Amz-") + +AWS_HEADERS = { + "x-amz-request-id": "REQ123", + "x-amz-id-2": "HOSTID456", + "Date": "Wed, 29 Jul 2026 12:31:07 GMT", +} + +# S3 rejecting a presigned URL whose signed window has elapsed +ACCESS_DENIED_EXPIRED_BODY = ( + b'\n' + b"AccessDeniedRequest has expired" + b"900" + b"2026-07-29T12:15:00Z" + b"2026-07-29T12:31:07Z" + b"REQ123HOSTID456" +) + +# S3 rejecting credentials that died before the URL's own expiry +EXPIRED_TOKEN_BODY = ( + b'\n' + b"ExpiredToken" + b"The provided token has expired." + b"REQ123HOSTID456" +) + +# S3 echoes the canonical request - and with it the signed query string - when the +# signature does not match, which is why the raw body must never be logged +SIGNATURE_MISMATCH_BODY = ( + b'\n' + b"SignatureDoesNotMatch" + b"The request signature we calculated does not match the signature " + b"you provided. Check your key and signing method." + b"CREDVALUE" + b"AWS4-HMAC-SHA256\n20260729T120000Z\n" + b"PUT\n/ws/int/key\n" + b"X-Amz-Credential=CREDVALUE&X-Amz-Security-Token=TOKENVALUE" + b"&X-Amz-Signature=SIGVALUE\nhost:bucket.s3.eu-west-1.amazonaws.com\n" + b"" + b"REQ123HOSTID456" +) + +# A gateway that answered instead of S3 and echoed the request line back. The +# field allowlist does not apply to a body that is not an S3 error document, so +# redaction is the only thing keeping the signed query string out of the snippet. +PROXY_ECHO_BODY = ( + b"502 Bad Gateway\n" + b"

502 Bad Gateway

\n" + b"

Upstream failed for request: PUT " + b"https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" + b"?X-Amz-Algorithm=AWS4-HMAC-SHA256" + b"&X-Amz-Credential=CREDVALUE%2F20260729%2Feu-west-1%2Fs3%2Faws4_request" + b"&X-Amz-Date=20260729T120000Z&X-Amz-Expires=900" + b"&X-Amz-Security-Token=TOKENVALUE&X-Amz-Signature=SIGVALUE

\n" + b"" +) + +# Every LogRecord attribute that logging refuses to let `extra` overwrite +RESERVED_LOGRECORD_ATTRS = set( + logging.LogRecord("", 0, "", 0, "", None, None).__dict__ +) | {"message", "asctime"} + + +def _s3_response(status_code, body=b"", headers=None): + """Build a stand-in for a requests.Response from the object store.""" + return mock.Mock(status_code=status_code, content=body, headers=headers or {}) + + +def _upload(url=PRESIGNED_URL, issued_at=None): + """Build an upload handle, issued now unless told otherwise.""" + return SqlCacheUpload( + url=url, issued_at=time.monotonic() if issued_at is None else issued_at + ) + + +def _cache_hit(download_url=PRESIGNED_URL): + """The webapp's answer when the query has a cached result waiting.""" + return { + "result": "cacheHit", + "downloadUrl": download_url, + "cacheCreatedAt": "2022-01-01 00:00:00", + } + + +def _logged_strings(mock_logger): + """Every string a logger.error call would hand to the error pipeline.""" + strings = [] + for call in mock_logger.error.call_args_list: + strings.extend(str(arg) for arg in call.args) + extra = call.kwargs.get("extra", {}) + strings.extend(str(key) for key in extra) + strings.extend(str(value) for value in extra.values()) + return strings + + +def _collect_logged_extras(): + """Run every failure path in the module and collect the extras it logs. + + Gathered in one place so the reserved-key and serializability guards cover all + of them, rather than whichever path a single test happened to exercise. + """ + dataframe = pd.DataFrame({"a": [1, 2, 3]}) + connection_error = requests.exceptions.ConnectionError(URLLIB3_ERROR_MESSAGE) + + with patch("deepnote_toolkit.sql.sql_caching.logger") as mock_logger: + with patch("deepnote_toolkit.sql.sql_caching.requests.put") as mock_put: + # the object store rejected the upload + mock_put.return_value = _s3_response( + 403, ACCESS_DENIED_EXPIRED_BODY, AWS_HEADERS + ) + upload_sql_cache(dataframe, _upload()) + + # the webapp answered with a null upload URL + upload_sql_cache(dataframe, SqlCacheUpload(url=None, issued_at=0.0)) + + # the upload request never completed + mock_put.side_effect = connection_error + upload_sql_cache(dataframe, _upload()) + + # the dataframe could not be serialized, so nothing was sent + unserializable = mock.Mock() + unserializable.to_parquet.side_effect = ValueError("column customer_email") + upload_sql_cache(unserializable, _upload()) + + with patch( + "deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp" + ) as mock_cache_info: + mock_cache_info.return_value = _cache_hit() + with patch("deepnote_toolkit.sql.sql_caching.requests.get") as mock_get: + # the object store rejected the download + mock_get.return_value = _s3_response( + 403, EXPIRED_TOKEN_BODY, AWS_HEADERS + ) + get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + # the download request never completed + mock_get.side_effect = connection_error + get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + # the webapp request itself failed + mock_cache_info.side_effect = connection_error + get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + # the webapp answered the cache info request with a non-200 + with ( + patch("deepnote_toolkit.sql.sql_caching.requests.get") as mock_get, + patch( + "deepnote_toolkit.sql.sql_caching.get_absolute_userpod_api_url" + ) as mock_url, + patch( + "deepnote_toolkit.sql.sql_caching.get_project_auth_headers" + ) as mock_headers, + ): + mock_url.return_value = ( + "http://localhost:19456/userpod-api/p1/integrations/123/sql-cache" + "?sqlCacheKey=abc&sqlCacheMode=read" + ) + mock_headers.return_value = {} + mock_get.return_value = _s3_response(503, b"upstream unavailable") + _request_cache_info_from_webapp("abc", "123", "read") + + return [call.kwargs["extra"] for call in mock_logger.error.call_args_list] + class TestGenerateCacheKey(unittest.TestCase): def test_empty_params_returns_valid_result(self): @@ -83,7 +284,7 @@ def test_cache_not_supported_for_query( mock_is_single_select_query.return_value = False - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) @@ -91,8 +292,9 @@ def test_cache_not_supported_for_query( {"status": "cache_not_supported_for_query"} ) self.assertIsNone(result_df) - self.assertIsNone(upload_url) + self.assertIsNone(upload) + @patch("deepnote_toolkit.sql.sql_caching.logger") @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @patch("deepnote_toolkit.sql.sql_caching.output_sql_metadata") @@ -101,6 +303,7 @@ def test_failed_to_request_cache_info( mock_output_sql_metadata, mock_request_cache_info_from_webapp, mock_is_single_select_query, + mock_logger, ): query = "SELECT * FROM users" bind_params = {} @@ -113,21 +316,23 @@ def test_failed_to_request_cache_info( "Failed to request cache info" ) - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) mock_output_sql_metadata.assert_not_called() self.assertIsNone(result_df) - self.assertIsNone(upload_url) + self.assertIsNone(upload) @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @patch("deepnote_toolkit.sql.sql_caching.output_sql_metadata") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") @patch("pandas.read_parquet") def test_read_from_cache_success( self, mock_read_parquet, + mock_get, mock_output_sql_metadata, mock_request_cache_info_from_webapp, mock_is_single_select_query, @@ -145,9 +350,10 @@ def test_read_from_cache_success( mock_is_single_select_query.return_value = True mock_request_cache_info_from_webapp.return_value = cache_info + mock_get.return_value = _s3_response(200, b"parquet-bytes") mock_read_parquet.return_value = pd.DataFrame() - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) @@ -161,17 +367,19 @@ def test_read_from_cache_success( } ) self.assertIsInstance(result_df, pd.DataFrame) - self.assertIsNone(upload_url) + self.assertIsNone(upload) @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @patch("deepnote_toolkit.sql.sql_caching.output_sql_metadata") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") @patch("pandas.read_parquet") @patch("pandas.read_pickle") def test_fallback_to_pickle_format( self, mock_read_pickle, mock_read_parquet, + mock_get, mock_output_sql_metadata, mock_request_cache_info_from_webapp, mock_is_single_select_query, @@ -189,10 +397,11 @@ def test_fallback_to_pickle_format( mock_is_single_select_query.return_value = True mock_request_cache_info_from_webapp.return_value = cache_info + mock_get.return_value = _s3_response(200, b"pickle-bytes") mock_read_parquet.side_effect = ArrowInvalid mock_read_pickle.return_value = pd.DataFrame() - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, @@ -210,18 +419,22 @@ def test_fallback_to_pickle_format( } ) self.assertIsInstance(result_df, pd.DataFrame) - self.assertIsNone(upload_url) + self.assertIsNone(upload) + @patch("deepnote_toolkit.sql.sql_caching.logger") @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @patch("deepnote_toolkit.sql.sql_caching.output_sql_metadata") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") @patch("pandas.read_parquet") def test_failed_to_download_from_cache( self, mock_read_parquet, + mock_get, mock_output_sql_metadata, mock_request_cache_info_from_webapp, mock_is_single_select_query, + mock_logger, ): query = "SELECT * FROM users" bind_params = {} @@ -236,14 +449,15 @@ def test_failed_to_download_from_cache( mock_is_single_select_query.return_value = True mock_request_cache_info_from_webapp.return_value = cache_info + mock_get.return_value = _s3_response(200, b"parquet-bytes") mock_read_parquet.side_effect = Exception("Failed to download from cache") - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) self.assertIsNone(result_df) - self.assertIsNone(upload_url) + self.assertIsNone(upload) @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @@ -263,12 +477,14 @@ def test_cache_miss( mock_is_single_select_query.return_value = True mock_request_cache_info_from_webapp.return_value = cache_info - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) self.assertIsNone(result_df) - self.assertEqual(upload_url, cache_info["uploadUrl"]) + self.assertIsInstance(upload, SqlCacheUpload) + self.assertEqual(upload.url, cache_info["uploadUrl"]) + self.assertIsInstance(upload.issued_at, float) @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @@ -288,12 +504,14 @@ def test_always_write( mock_is_single_select_query.return_value = True mock_request_cache_info_from_webapp.return_value = cache_info - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) self.assertIsNone(result_df) - self.assertEqual(upload_url, cache_info["uploadUrl"]) + self.assertIsInstance(upload, SqlCacheUpload) + self.assertEqual(upload.url, cache_info["uploadUrl"]) + self.assertIsInstance(upload.issued_at, float) @patch("deepnote_toolkit.sql.sql_caching.is_single_select_query") @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") @@ -309,69 +527,436 @@ def test_no_cache_info( mock_is_single_select_query.return_value = True mock_request_cache_info_from_webapp.return_value = None - result_df, upload_url = get_sql_cache( + result_df, upload = get_sql_cache( query, bind_params, integration_id, sql_cache_mode, return_variable_type ) self.assertIsNone(result_df) - self.assertIsNone(upload_url) + self.assertIsNone(upload) + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") @patch("pandas.read_parquet") @patch("pandas.read_pickle") def test_read_from_cache_error_doesnt_raise( - self, mock_read_pickle, mock_read_parquet + self, + mock_read_pickle, + mock_read_parquet, + mock_get, + mock_cache_info, + mock_logger, ): + mock_cache_info.return_value = _cache_hit("https://example.com/cache") + mock_get.return_value = _s3_response(200, b"not-a-dataframe") mock_read_parquet.side_effect = ArrowInvalid mock_read_pickle.side_effect = Exception("Error reading pickle") - query = "SELECT * FROM users" - bind_params = {} - integration_id = "123" - sql_cache_mode = "read" - return_variable_type = "dataframe" + result_df, upload = get_sql_cache(QUERY, {}, "123", "read", "dataframe") - result_df, upload_url = get_sql_cache( - query, bind_params, integration_id, sql_cache_mode, return_variable_type + # both read attempts ran, and the failure of the second was swallowed + mock_read_parquet.assert_called_once() + mock_read_pickle.assert_called_once() + self.assertIsNone(result_df) + self.assertIsNone(upload) + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") + def test_download_http_error_logs_s3_diagnostics( + self, mock_get, mock_cache_info, mock_logger + ): + mock_cache_info.return_value = _cache_hit() + mock_get.return_value = _s3_response(403, EXPIRED_TOKEN_BODY, AWS_HEADERS) + + result_df, upload = get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + self.assertIsNone(result_df) + self.assertIsNone(upload) + mock_logger.error.assert_called_once() + self.assertEqual( + mock_logger.error.call_args.args, + ("Failed to download dataframe from cache",), + ) + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["sql_caching_cause"], "failed_to_download_from_cache") + self.assertEqual(extra["s3_error_code"], "ExpiredToken") + self.assertEqual(extra["status_code"], 403) + self.assertEqual(extra["aws_request_id"], "REQ123") + self.assertEqual(extra["object_path"], "/ws/int/key") + + @parameterized.expand( + [ + # the field allowlist is what keeps this one clean + ("signature_mismatch_body", SIGNATURE_MISMATCH_BODY), + # no allowlist applies here, so only redaction keeps it clean + ("proxy_echoing_request_url", PROXY_ECHO_BODY), + ] + ) + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") + def test_download_http_error_logs_no_query_string( + self, _name, body, mock_get, mock_cache_info, mock_logger + ): + mock_cache_info.return_value = _cache_hit() + mock_get.return_value = _s3_response(403, body, AWS_HEADERS) + + get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + mock_logger.error.assert_called_once() + for logged in _logged_strings(mock_logger): + for secret in SECRETS: + self.assertNotIn(secret, logged) + + @patch("deepnote_toolkit.sql.sql_caching.output_sql_metadata") + @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") + def test_download_reads_parquet_from_fetched_bytes( + self, mock_get, mock_cache_info, mock_output_sql_metadata + ): + dataframe = pd.DataFrame({"a": [1, 2, 3]}) + buffer = BytesIO() + dataframe.to_parquet(buffer) + + mock_cache_info.return_value = _cache_hit() + mock_get.return_value = _s3_response(200, buffer.getvalue()) + + result_df, _ = get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + pd.testing.assert_frame_equal(result_df, dataframe) + # the object was fetched here rather than by handing the URL to pandas + mock_get.assert_called_once() + self.assertEqual(mock_get.call_args.args[0], PRESIGNED_URL) + + @patch("deepnote_toolkit.sql.sql_caching.output_sql_metadata") + @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") + def test_download_falls_back_to_pickle_from_same_bytes( + self, mock_get, mock_cache_info, mock_output_sql_metadata + ): + dataframe = pd.DataFrame({"a": [1, 2, 3]}) + buffer = BytesIO() + dataframe.to_pickle(buffer) + + mock_cache_info.return_value = _cache_hit() + mock_get.return_value = _s3_response(200, buffer.getvalue()) + + result_df, _ = get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + # pins the seek(0) that the failed parquet read makes necessary + pd.testing.assert_frame_equal(result_df, dataframe) + # and the single fetch that replaced one download per format attempt + mock_get.assert_called_once() + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching._request_cache_info_from_webapp") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") + def test_download_network_error_is_swallowed( + self, mock_get, mock_cache_info, mock_logger + ): + mock_cache_info.return_value = _cache_hit() + mock_get.side_effect = requests.exceptions.ConnectionError( + URLLIB3_ERROR_MESSAGE ) + result_df, upload = get_sql_cache(QUERY, {}, "123", "read", "dataframe") + self.assertIsNone(result_df) - self.assertIsNone(upload_url) + self.assertIsNone(upload) + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["error_type"], "ConnectionError") + for logged in _logged_strings(mock_logger): + for secret in SECRETS: + self.assertNotIn(secret, logged) + + +class TestRequestCacheInfoFromWebapp(unittest.TestCase): + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.get_project_auth_headers") + @patch("deepnote_toolkit.sql.sql_caching.get_absolute_userpod_api_url") + @patch("deepnote_toolkit.sql.sql_caching.requests.get") + def test_non_200_logs_constant_message_with_status_in_extra( + self, mock_get, mock_url, mock_headers, mock_logger + ): + mock_url.return_value = ( + "http://localhost:19456/userpod-api/p1/integrations/123/sql-cache" + "?sqlCacheKey=abc&sqlCacheMode=read" + ) + mock_headers.return_value = {} + mock_get.return_value = _s3_response(503, b"upstream unavailable") + + self.assertIsNone(_request_cache_info_from_webapp("abc", "123", "read")) + + self.assertEqual( + mock_logger.error.call_args.args, ("Failed to request cache info",) + ) + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["status_code"], 503) + self.assertEqual(extra["response_body_snippet"], "upstream unavailable") + self.assertEqual( + extra["cache_info_path"], + "/userpod-api/p1/integrations/123/sql-cache", + ) + # the query string of the request is not part of the entry + self.assertNotIn("sqlCacheKey", " ".join(_logged_strings(mock_logger))) + + +class TestRedactSensitive(unittest.TestCase): + def test_strips_query_string_from_urlopen_style_message(self): + # the shape urllib3 actually produces: the URL is path-only + redacted = _redact_sensitive(URLLIB3_ERROR_MESSAGE) + + self.assertIn("bucket.s3.eu-west-1.amazonaws.com", redacted) + self.assertIn("/ws/int/key?", redacted) + for secret in SECRETS: + self.assertNotIn(secret, redacted) + + def test_query_string_strip_alone_redacts_urllib3_message(self): + """The primary defence must hold without help from the parameter backstop.""" + stripped = _URL_QUERY_PATTERN.sub(r"\1?", URLLIB3_ERROR_MESSAGE) + + for secret in SECRETS: + self.assertNotIn(secret, stripped) + + def test_blanks_aws_params_outside_a_url(self): + redacted = _redact_sensitive( + "X-Amz-Credential=CREDVALUE&X-Amz-Security-Token=TOKENVALUE" + "&X-Amz-Signature=SIGVALUE" + ) + + self.assertEqual( + redacted, + "X-Amz-Credential=&X-Amz-Security-Token=" + "&X-Amz-Signature=", + ) + + @parameterized.expand( + [ + ("question_in_prose", "Is this ok? Yes it is."), + ("bare_question", "what? nothing"), + ("plain_sentence", "The provided token has expired."), + ] + ) + def test_leaves_ordinary_text_unchanged(self, _, text): + self.assertEqual(_redact_sensitive(text), text) + + +class TestDescribeException(unittest.TestCase): + def test_large_message_is_bounded_and_redacted_in_bounded_time(self): + """The message is cut before redacting, not after. + + _URL_QUERY_PATTERN backtracks over every start position in a long run of + non-separator characters, so redacting an untruncated exception message is + quadratic - 64k characters took 8.7s, charged to the user's cell. + """ + message = URLLIB3_ERROR_MESSAGE + "&padding=" + "A" * 64_000 + + started = time.monotonic() + described = _describe_exception(ValueError(message)) + elapsed = time.monotonic() - started + + self.assertEqual(described["error_type"], "ValueError") + self.assertLessEqual(len(described["error_message"]), 500) + # cutting the query string short is still safe: the pattern runs to the + # end of the string, so the remainder is stripped either way + self.assertIn("/ws/int/key?", described["error_message"]) + for secret in SECRETS: + self.assertNotIn(secret, described["error_message"]) + self.assertLess(elapsed, 2.0) + + +class TestDescribeS3Error(unittest.TestCase): + def test_extracts_code_and_message_from_xml(self): + diagnostics = _describe_s3_error( + _s3_response(403, ACCESS_DENIED_EXPIRED_BODY, AWS_HEADERS) + ) + + self.assertEqual(diagnostics["status_code"], 403) + self.assertEqual(diagnostics["s3_error_code"], "AccessDenied") + self.assertIn("Request has expired", diagnostics["s3_error_message"]) + # AWS's own account of when the URL died and what time it was + self.assertEqual(diagnostics["s3_expires"], "2026-07-29T12:15:00Z") + self.assertEqual(diagnostics["s3_server_time"], "2026-07-29T12:31:07Z") + + def test_extracts_expired_token_body(self): + diagnostics = _describe_s3_error(_s3_response(403, EXPIRED_TOKEN_BODY)) + + self.assertEqual(diagnostics["s3_error_code"], "ExpiredToken") + self.assertEqual( + diagnostics["s3_error_message"], "The provided token has expired." + ) + + def test_captures_aws_request_headers(self): + diagnostics = _describe_s3_error( + _s3_response(403, EXPIRED_TOKEN_BODY, AWS_HEADERS) + ) + + self.assertEqual(diagnostics["aws_request_id"], "REQ123") + self.assertEqual(diagnostics["aws_host_id"], "HOSTID456") + self.assertEqual(diagnostics["aws_date"], "Wed, 29 Jul 2026 12:31:07 GMT") + + def test_signature_mismatch_body_surfaces_only_code_and_message(self): + """Pins the field allowlist: and friends are never read. + + Redaction is not what protects this body - the elements carrying the signed + query string are simply not among the four that get extracted. + """ + diagnostics = _describe_s3_error( + _s3_response(403, SIGNATURE_MISMATCH_BODY, AWS_HEADERS) + ) + + self.assertEqual(diagnostics["s3_error_code"], "SignatureDoesNotMatch") + # an S3 error document was recognised, so no free-text snippet is kept + self.assertNotIn("response_body_snippet", diagnostics) + for value in diagnostics.values(): + for secret in SECRETS: + self.assertNotIn(secret, str(value)) + + def test_proxy_body_echoing_request_url_is_redacted(self): + """Pins redaction: the allowlist cannot help once the body is not S3's.""" + diagnostics = _describe_s3_error(_s3_response(502, PROXY_ECHO_BODY)) + + self.assertIsNone(diagnostics["s3_error_code"]) + snippet = diagnostics["response_body_snippet"] + # the request line survives, the credentials on it do not + self.assertIn("502 Bad Gateway", snippet) + self.assertIn("/ws/int/key?", snippet) + for secret in SECRETS: + self.assertNotIn(secret, snippet) + + def test_non_xml_body_yields_snippet_without_code(self): + diagnostics = _describe_s3_error( + _s3_response(502, b"502 Bad Gateway") + ) + + self.assertIsNone(diagnostics["s3_error_code"]) + self.assertIsNone(diagnostics["aws_request_id"]) + self.assertIn("502 Bad Gateway", diagnostics["response_body_snippet"]) + self.assertLessEqual(len(diagnostics["response_body_snippet"]), 500) + + def test_oversized_body_is_bounded(self): + body = ( + b"AccessDenied" + + b"x" * 1000 + + b"" + + b"y" * 10_000 + + b"" + ) + + diagnostics = _describe_s3_error(_s3_response(403, body)) + + self.assertEqual(diagnostics["s3_error_code"], "AccessDenied") + for value in diagnostics.values(): + if isinstance(value, str): + self.assertLessEqual(len(value), 500) + + +class TestDescribePresignedUrl(unittest.TestCase): + def test_returns_path_and_expiry(self): + described = _describe_presigned_url(PRESIGNED_URL) + + self.assertEqual(described["object_host"], "bucket.s3.eu-west-1.amazonaws.com") + self.assertEqual(described["object_path"], "/ws/int/key") + self.assertEqual(described["url_expires_in"], 900) + + def test_never_returns_query_string(self): + described = _describe_presigned_url(PRESIGNED_URL) + + for value in described.values(): + for secret in SECRETS: + self.assertNotIn(secret, str(value)) + + def test_missing_expires_yields_none(self): + described = _describe_presigned_url("https://example.com/x") + + self.assertEqual(described["object_path"], "/x") + self.assertIsNone(described["url_expires_in"]) + + def test_non_numeric_expires_yields_none(self): + described = _describe_presigned_url("https://example.com/x?X-Amz-Expires=abc") + + self.assertIsNone(described["url_expires_in"]) + + @parameterized.expand( + [ + ("unterminated_ipv6", "https://[::1"), + ("empty", ""), + ("not_a_url", "not a url"), + # urlsplit answers these with bytes fields instead of raising + ("none", None), + ("bytes", b"/ws/int/key"), + ("dict", {"url": "https://example.com/x"}), + ] + ) + def test_malformed_url_does_not_raise_or_leak(self, _, url): + described = _describe_presigned_url(url) + + self.assertIsNone(described["url_expires_in"]) + for value in described.values(): + for secret in SECRETS: + self.assertNotIn(secret, str(value)) + # a bytes value in either would silently discard the entire error report + json.dumps(described) + json.dumps(_safe_url_path(url)) class TestUploadSqlCache(unittest.TestCase): + @patch("deepnote_toolkit.sql.sql_caching.logger") @patch("deepnote_toolkit.sql.sql_caching.requests.put") - def test_upload_parquet_success(self, mock_put): - mock_put.return_value = mock.Mock(raise_for_status=mock.Mock()) + def test_upload_parquet_success(self, mock_put, mock_logger): + mock_put.return_value = mock.Mock(status_code=200) df = pd.DataFrame({"a": [1, 2, 3]}) - upload_sql_cache(df, "https://example.com/upload") + upload_sql_cache(df, _upload("https://example.com/upload")) mock_put.assert_called_once() args, _ = mock_put.call_args self.assertEqual(args[0], "https://example.com/upload") + mock_logger.error.assert_not_called() @patch("deepnote_toolkit.sql.sql_caching.requests.put") def test_overflow_error_falls_back_to_pickle(self, mock_put): """Large Python int triggers OverflowError in to_parquet, upload succeeds via pickle.""" uploaded_bytes = None - def capture_put(_url, data): + def capture_put(_url, data, **_kwargs): nonlocal uploaded_bytes uploaded_bytes = data.read() - return mock.Mock(raise_for_status=mock.Mock()) + return mock.Mock(status_code=200) mock_put.side_effect = capture_put df = pd.DataFrame({"x": pd.array([2**100, 1], dtype=object)}) - upload_sql_cache(df, "https://example.com/upload") + upload_sql_cache(df, _upload("https://example.com/upload")) roundtripped = pd.read_pickle(pd.io.common.BytesIO(uploaded_bytes)) pd.testing.assert_frame_equal(roundtripped, df) + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_arrow_failure_still_falls_back_to_pickle_and_uploads(self, mock_put): + """The narrowed serialization handler must not swallow the pickle fallback.""" + uploaded_bytes = None + + def capture_put(_url, data, **_kwargs): + nonlocal uploaded_bytes + uploaded_bytes = data.read() + return mock.Mock(status_code=200) + + mock_put.side_effect = capture_put + # nested dicts of mixed shape are not representable in parquet + df = pd.DataFrame({"x": [{"a": 1}, {"a": "two"}]}) + + upload_sql_cache(df, _upload("https://example.com/upload")) + + mock_put.assert_called_once() + roundtripped = pd.read_pickle(BytesIO(uploaded_bytes)) + pd.testing.assert_frame_equal(roundtripped, df) + @patch("deepnote_toolkit.sql.sql_caching.requests.put") def test_pickle_fallback_truncates_partial_parquet_bytes(self, mock_put): """When to_parquet writes partial bytes before failing, truncate clears them.""" - mock_put.return_value = mock.Mock(raise_for_status=mock.Mock()) + mock_put.return_value = mock.Mock(status_code=200) def write_garbage_then_overflow(f, **_kwargs): f.write(b"partial parquet data") @@ -390,7 +975,204 @@ def capture_file_state(f, **_kwargs): df.to_parquet.side_effect = write_garbage_then_overflow df.to_pickle.side_effect = capture_file_state - upload_sql_cache(df, "https://example.com/upload") + upload_sql_cache(df, _upload("https://example.com/upload")) self.assertEqual(pickle_pos, 0, "file should be at position 0") self.assertEqual(pickle_size, 0, "file should be empty after truncate") + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_http_error_logs_s3_diagnostics(self, mock_put, mock_logger): + mock_put.return_value = _s3_response( + 403, ACCESS_DENIED_EXPIRED_BODY, AWS_HEADERS + ) + + upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) + + mock_logger.error.assert_called_once() + # a single constant message, with no format placeholders and no args + self.assertEqual( + mock_logger.error.call_args.args, ("Failed to upload SQL cache",) + ) + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["sql_caching_cause"], "failed_to_upload_to_cache") + self.assertEqual(extra["s3_error_code"], "AccessDenied") + self.assertEqual(extra["s3_error_message"], "Request has expired") + self.assertEqual(extra["status_code"], 403) + self.assertEqual(extra["aws_request_id"], "REQ123") + self.assertEqual(extra["aws_host_id"], "HOSTID456") + self.assertEqual(extra["object_path"], "/ws/int/key") + self.assertEqual(extra["url_expires_in"], 900) + self.assertGreaterEqual(extra["seconds_since_url_issued"], 0) + + @parameterized.expand( + [ + # the field allowlist is what keeps this one clean + ("signature_mismatch_body", SIGNATURE_MISMATCH_BODY), + # no allowlist applies here, so only redaction keeps it clean + ("proxy_echoing_request_url", PROXY_ECHO_BODY), + ] + ) + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_http_error_logs_no_presigned_query_string( + self, _name, body, mock_put, mock_logger + ): + mock_put.return_value = _s3_response(403, body, AWS_HEADERS) + + upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) + + mock_logger.error.assert_called_once() + for logged in _logged_strings(mock_logger): + for secret in SECRETS: + self.assertNotIn(secret, logged) + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_network_error_logs_no_presigned_url(self, mock_put, mock_logger): + mock_put.side_effect = requests.exceptions.ConnectionError( + URLLIB3_ERROR_MESSAGE + ) + + upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) + + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["sql_caching_cause"], "failed_to_upload_to_cache") + self.assertEqual(extra["error_type"], "ConnectionError") + for logged in _logged_strings(mock_logger): + for secret in SECRETS: + self.assertNotIn(secret, logged) + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_message_is_constant_across_different_failures(self, mock_put, mock_logger): + df = pd.DataFrame({"a": [1, 2, 3]}) + other_url = PRESIGNED_URL.replace("/ws/int/key", "/other/int/key2") + + mock_put.return_value = _s3_response(403, ACCESS_DENIED_EXPIRED_BODY) + upload_sql_cache(df, _upload()) + mock_put.return_value = _s3_response(500, b"Slow") + upload_sql_cache(df, _upload(other_url)) + mock_put.side_effect = requests.exceptions.ConnectionError( + URLLIB3_ERROR_MESSAGE + ) + upload_sql_cache(df, _upload()) + upload_sql_cache(df, _upload(other_url)) + + messages = {call.args for call in mock_logger.error.call_args_list} + self.assertEqual(messages, {("Failed to upload SQL cache",)}) + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_serialization_failure_uses_distinct_cause(self, mock_put, mock_logger): + df = mock.Mock() + df.to_parquet.side_effect = ValueError("Conversion failed for column secret") + df.to_pickle.side_effect = ValueError("Conversion failed for column secret") + + upload_sql_cache(df, _upload()) + + mock_put.assert_not_called() + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["sql_caching_cause"], "failed_to_serialize_cache") + self.assertEqual(extra["error_type"], "ValueError") + # the message names the user's columns, so it is deliberately not logged + self.assertNotIn("error_message", extra) + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_seconds_since_url_issued_reflects_elapsed_time( + self, mock_put, mock_logger + ): + mock_put.return_value = _s3_response(403, ACCESS_DENIED_EXPIRED_BODY) + + upload_sql_cache( + pd.DataFrame({"a": [1, 2, 3]}), + _upload(issued_at=time.monotonic() - 930), + ) + + extra = mock_logger.error.call_args.kwargs["extra"] + # the entry alone shows the URL was used after its window closed + self.assertGreaterEqual(extra["seconds_since_url_issued"], 930) + self.assertEqual(extra["url_expires_in"], 900) + + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_non_numeric_issued_at_does_not_raise(self, mock_put, mock_logger): + """The final logger.error is outside the try, so nothing in it may raise.""" + mock_put.return_value = _s3_response(403, ACCESS_DENIED_EXPIRED_BODY) + + upload_sql_cache( + pd.DataFrame({"a": [1, 2, 3]}), + SqlCacheUpload(url=PRESIGNED_URL, issued_at="not-a-number"), + ) + + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertIsNone(extra["seconds_since_url_issued"]) + + @parameterized.expand( + [ + ("http_403", 403, ACCESS_DENIED_EXPIRED_BODY, None, False), + ( + "http_500", + 500, + b"InternalError", + None, + False, + ), + ("connection_error", None, None, "ConnectionError", False), + ("timeout", None, None, "Timeout", False), + ("serialization_error", 200, b"", None, True), + ] + ) + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_upload_failure_never_raises( + self, + _name, + status_code, + body, + put_exception, + fail_serialization, + mock_put, + mock_logger, + ): + if put_exception is not None: + mock_put.side_effect = getattr(requests.exceptions, put_exception)( + URLLIB3_ERROR_MESSAGE + ) + else: + mock_put.return_value = _s3_response(status_code, body) + + if fail_serialization: + dataframe = mock.Mock() + dataframe.to_parquet.side_effect = ValueError("cannot serialize") + else: + dataframe = pd.DataFrame({"a": [1, 2, 3]}) + + self.assertIsNone(upload_sql_cache(dataframe, _upload())) + + +class TestLoggedExtras(unittest.TestCase): + """Guards that apply to every extra dict this module logs.""" + + def setUp(self): + self.extras = _collect_logged_extras() + + def test_every_failure_path_logs_an_extra(self): + # upload http, upload with a null url, upload network, serialization, + # download http, download network, cache info exception, cache info non-200 + self.assertEqual(len(self.extras), 8) + for extra in self.extras: + self.assertIn("sql_caching_cause", extra) + + def test_extra_keys_avoid_reserved_logrecord_attributes(self): + """A reserved key makes logger.error raise into the user's query.""" + for extra in self.extras: + self.assertEqual(set(extra) & RESERVED_LOGRECORD_ATTRS, set()) + + def test_extra_is_json_serializable(self): + """A non-serializable value silently discards the whole error report.""" + for extra in self.extras: + for value in extra.values(): + self.assertIsInstance(value, (str, int, float, bool, type(None))) + json.dumps(extra)