From c9fa363ba4d4adfdfe2df4a63fa92ca65b1206d1 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 29 Jul 2026 15:42:33 +0000 Subject: [PATCH 1/5] 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) From ece6d20cee8452f2e7dcc636d57099a0c3321c7c Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 13:25:30 +0000 Subject: [PATCH 2/5] fix: bound cache transfers, fix chunked error reads and URL-age timing Applies review findings on the SQL cache diagnostics: - Bound the read phase of object store requests. The previous comment held that a read timeout would abort slow but healthy transfers; requests applies it per received chunk, not to the transfer as a whole, so it only ever aborts a connection that has gone silent. - Read a failed download's body off the wire instead of buffering it whole, and accumulate chunks up to the cap. A chunked response answers with one HTTP chunk per read however large a size is asked for, so reading once lost and entirely on a small-chunked error body. - Measure the presigned URL's age when the PUT starts rather than when the failure is logged. S3 checks a signature on arrival, so folding the transfer into the age made a slow upload read back as an expired URL. Transfer time is now reported separately as upload_duration_seconds. - Cut an over-encoded query string off object_path before logging it. urlsplit only splits on a literal '?', so a URL that encoded it kept its signing parameters in the path. - Drop comments that restate the code, docstrings that restate signatures, a guard against an input that cannot occur, and tests that cannot fail unless another test also fails. The response mock now yields chunked bodies and empties on block exit, both of which previously hid real regressions from the suite. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_018cayDdcs1iuui8Pgh5npKm --- deepnote_toolkit/sql/sql_caching.py | 175 +++++++++------------- tests/unit/test_sql_caching.py | 223 ++++++++++++++++++---------- 2 files changed, 214 insertions(+), 184 deletions(-) diff --git a/deepnote_toolkit/sql/sql_caching.py b/deepnote_toolkit/sql/sql_caching.py index 3f325b6..b23486e 100644 --- a/deepnote_toolkit/sql/sql_caching.py +++ b/deepnote_toolkit/sql/sql_caching.py @@ -20,23 +20,17 @@ # 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) +# The read timeout bounds the gap between two received chunks, not the transfer as +# a whole, so it cannot abort a large but healthy download. +_OBJECT_STORE_TIMEOUT: tuple[int, int] = (5, 60) -# 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. +# Matches the scheme-less, path-only form that urllib3 puts in its connection +# error messages as well as a whole URL _URL_QUERY_PATTERN = re.compile( r"((?:https?://)?[^\s?\"'<>]*/[^\s?\"'<>]*)\?[^\s\"'<>]*" ) @@ -46,53 +40,37 @@ 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. +# urlsplit only splits on a literal '?', so an over-encoded URL keeps its signed +# query string in the path +_ENCODED_QUERY_SEPARATOR = re.compile("%3F", re.IGNORECASE) +# 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. - """ + """A presigned upload URL together with the moment it was issued.""" 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. + 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. - """ + """Store log-safe diagnostics already taken from the failed 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. - """ + """Strip credential-bearing material from text destined for a log.""" redacted = _URL_QUERY_PATTERN.sub(r"\1?", text) return _AWS_CREDENTIAL_PARAM_PATTERN.sub(r"\1", redacted) @@ -108,27 +86,17 @@ def _to_int_or_none(value: Optional[str]) -> Optional[int]: 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: +def _seconds_between(start: Optional[float], end: Optional[float]) -> Optional[float]: + """Monotonic seconds from start to end, or None when either was not taken.""" + if start is None or end is None: return None + return round(end - start, 1) -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. - """ +def _safe_url_path(url: str) -> Optional[str]: + """The bounded path of a URL, or None 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 None try: @@ -138,12 +106,7 @@ def _safe_url_path(url: str) -> Optional[str]: 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. - """ + """Object path and declared expiry, without the credential-bearing query string.""" 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 @@ -152,43 +115,49 @@ def _describe_presigned_url(url: str) -> dict[str, Any]: try: parts = urlsplit(url) expires_values = parse_qs(parts.query).get("X-Amz-Expires") + path = _ENCODED_QUERY_SEPARATOR.split(parts.path, maxsplit=1)[0] return { # hostname rather than netloc, which would carry any user:pass@ userinfo "object_host": parts.hostname, - "object_path": parts.path[:_MAX_OBJECT_PATH_CHARS], + # cut for an encoded '?', redacted for separators urlsplit ignores + "object_path": _redact_sensitive(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. + Taken off the wire rather than sliced off response.content, which would + download the whole body first, and accumulated rather than read once, because + a chunked response answers with one HTTP chunk per read however large a size + is asked for. """ try: - return response.content[:_MAX_ERROR_BODY_BYTES].decode( - "utf-8", errors="replace" - ) + chunks = [] + length = 0 + for chunk in response.iter_content(_MAX_ERROR_BODY_BYTES): + chunks.append(chunk) + length += len(chunk) + if length >= _MAX_ERROR_BODY_BYTES: + break + + prefix = b"".join(chunks)[:_MAX_ERROR_BODY_BYTES] + return prefix.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. + Only the allowlisted fields are taken from the body, because the other + elements of an S3 error document (, , + ) echo the signed query string and with it the credentials. """ diagnostics: dict[str, Any] = { "status_code": response.status_code, @@ -221,8 +190,7 @@ def _describe_s3_error(response: requests.Response) -> dict[str, Any]: 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. + # Not an S3 error document - a proxy or gateway answered instead diagnostics["response_body_snippet"] = _redact_sensitive(body)[ :_MAX_ERROR_FIELD_CHARS ] @@ -236,9 +204,7 @@ def _describe_exception(exc: BaseException) -> dict[str, Any]: 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. + # start position, and an exception message has no bound of its own message = _redact_sensitive(str(exc)[:_MAX_RAW_EXCEPTION_CHARS]) return { "error_type": type(exc).__name__, @@ -281,8 +247,7 @@ 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 + # inside this round trip, which makes it a conservative upper bound requested_at = time.monotonic() cache_info = None @@ -360,20 +325,16 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: 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. """ + put_started_at: Optional[float] = None try: with tempfile.TemporaryFile() as temp_file: try: _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. + # Only the type is logged: serialization errors quote the user's + # column names logger.error( "Failed to upload SQL cache", extra={ @@ -384,7 +345,9 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: return temp_file.seek(0) - # PUT the file to the pre-signed s3 url + # S3 checks a presigned signature when the request arrives rather than + # when it finishes, so this is the age that decided whether it was valid + put_started_at = time.monotonic() response = requests.put( upload.url, data=temp_file, timeout=_OBJECT_STORE_TIMEOUT ) @@ -396,7 +359,6 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: # interpolates the whole presigned URL and discards S3's error body failure = _describe_s3_error(response) except Exception as exc: - # The request never completed, so there is no response to describe failure = _describe_exception(exc) logger.error( @@ -405,7 +367,14 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: "sql_caching_cause": "failed_to_upload_to_cache", **failure, **_describe_presigned_url(upload.url), - "seconds_since_url_issued": _seconds_since(upload.issued_at), + "seconds_since_url_issued": _seconds_between( + upload.issued_at, put_started_at + ), + # Separate field so "the URL was already old" stays distinguishable + # from "the transfer was slow" + "upload_duration_seconds": _seconds_between( + put_started_at, time.monotonic() + ), }, ) @@ -416,11 +385,18 @@ def _try_read_cache(download_url: str) -> pd.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)) + # Streamed so that a failed download costs a bounded prefix rather than the + # whole error body, and so the rest of it is dropped with the connection + with requests.get( + download_url, timeout=_OBJECT_STORE_TIMEOUT, stream=True + ) as response: + if response.status_code >= 400: + raise _SqlCacheHttpError(_describe_s3_error(response)) + + # Read inside the block: once it exits, the body reads back as empty + # rather than raising, which would cache-miss every hit in silence + buffer = BytesIO(response.content) - buffer = BytesIO(response.content) try: # Attempt to read as a parquet file return pd.read_parquet(buffer) @@ -431,8 +407,6 @@ def _try_read_cache(download_url: str) -> pd.DataFrame: # (see .to_pickle fallback in upload_sql_cache) pass - # 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) @@ -446,16 +420,7 @@ def _generate_cache_key(query, bind_params): 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. - """ + """The cache info for this query, or None when caching is off 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}" diff --git a/tests/unit/test_sql_caching.py b/tests/unit/test_sql_caching.py index 031fbfe..01395aa 100644 --- a/tests/unit/test_sql_caching.py +++ b/tests/unit/test_sql_caching.py @@ -28,8 +28,7 @@ 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 +# Signing 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" @@ -50,7 +49,6 @@ "(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 = { @@ -59,7 +57,6 @@ "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" @@ -77,8 +74,8 @@ 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 +# S3 echoes the canonical request - and with it the signed query string - when +# the signature does not match SIGNATURE_MISMATCH_BODY = ( b'\n' b"SignatureDoesNotMatch" @@ -93,9 +90,7 @@ 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. +# A gateway that answered instead of S3 and echoed the request line back PROXY_ECHO_BODY = ( b"502 Bad Gateway\n" b"

502 Bad Gateway

\n" @@ -108,7 +103,6 @@ b"" ) -# Every LogRecord attribute that logging refuses to let `extra` overwrite RESERVED_LOGRECORD_ATTRS = set( logging.LogRecord("", 0, "", 0, "", None, None).__dict__ ) | {"message", "asctime"} @@ -116,7 +110,18 @@ 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 {}) + response = mock.MagicMock( + status_code=status_code, content=body, headers=headers or {} + ) + response.__enter__.return_value = response + # a real streamed body reads back empty once the block exits, silently + response.__exit__.side_effect = lambda *_: setattr(response, "content", b"") + # a chunked response yields one HTTP chunk per read however large a size is + # asked for + response.iter_content.side_effect = lambda size: iter( + [body[i : i + 20] for i in range(0, len(body), 20)] + ) + return response def _upload(url=PRESIGNED_URL, issued_at=None): @@ -157,20 +162,16 @@ def _collect_logged_extras(): 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()) @@ -180,21 +181,17 @@ def _collect_logged_extras(): ) 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( @@ -649,6 +646,28 @@ def test_download_falls_back_to_pickle_from_same_bytes( # and the single fetch that replaced one download per format attempt mock_get.assert_called_once() + @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_request_is_bounded_and_streamed( + 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()) + + get_sql_cache(QUERY, {}, "123", "read", "dataframe") + + # an unbounded read phase hangs the user's cell on a socket gone silent + connect_timeout, read_timeout = mock_get.call_args.kwargs["timeout"] + self.assertIsNotNone(connect_timeout) + self.assertIsNotNone(read_timeout) + # and an unstreamed error body costs the size of that body + self.assertTrue(mock_get.call_args.kwargs["stream"]) + @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") @@ -698,7 +717,6 @@ def test_non_200_logs_constant_message_with_status_in_extra( 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))) @@ -807,7 +825,6 @@ def test_signature_mismatch_body_surfaces_only_code_and_message(self): ) 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: @@ -851,6 +868,22 @@ def test_oversized_body_is_bounded(self): if isinstance(value, str): self.assertLessEqual(len(value), 500) + def test_body_prefix_is_streamed_rather_than_buffered(self): + """Reading .content downloads the whole body just to keep 4 KB of it.""" + buffered = [] + response = mock.MagicMock(status_code=403, headers={}) + type(response).content = mock.PropertyMock( + side_effect=lambda: buffered.append("content") + ) + response.iter_content.side_effect = lambda size: iter( + [ACCESS_DENIED_EXPIRED_BODY[:size]] + ) + + diagnostics = _describe_s3_error(response) + + self.assertEqual(buffered, []) + self.assertEqual(diagnostics["s3_error_code"], "AccessDenied") + class TestDescribePresignedUrl(unittest.TestCase): def test_returns_path_and_expiry(self): @@ -860,13 +893,6 @@ def test_returns_path_and_expiry(self): 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") @@ -883,7 +909,6 @@ def test_non_numeric_expires_yields_none(self): ("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"}), @@ -900,6 +925,38 @@ def test_malformed_url_does_not_raise_or_leak(self, _, url): json.dumps(described) json.dumps(_safe_url_path(url)) + @parameterized.expand( + [ + ( + "separator_encoded", + "%3FX-Amz-Credential=CREDVALUE&X-Amz-Security-Token=TOKENVALUE" + "&X-Amz-Signature=SIGVALUE", + ), + ( + "separator_and_equals_encoded", + "%3FX-Amz-Credential%3DCREDVALUE&X-Amz-Security-Token%3DTOKENVALUE" + "&X-Amz-Signature%3DSIGVALUE", + ), + ( + "whole_query_encoded", + "%3FX-Amz-Credential%3DCREDVALUE%26X-Amz-Security-Token%3DTOKENVALUE" + "%26X-Amz-Signature%3DSIGVALUE", + ), + ("separator_is_a_semicolon", ";X-Amz-Credential=CREDVALUE"), + ("no_separator_at_all", "X-Amz-Signature=SIGVALUE"), + ] + ) + def test_over_encoded_url_does_not_leak_signing_params_via_path(self, _, suffix): + """urlsplit only splits on a literal '?', so the path carries the rest.""" + described = _describe_presigned_url( + "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" + suffix + ) + + # the parameter names may survive redaction, their values must not + for value in described.values(): + for secret in ("CREDVALUE", "TOKENVALUE", "SIGVALUE"): + self.assertNotIn(secret, str(value)) + class TestUploadSqlCache(unittest.TestCase): @patch("deepnote_toolkit.sql.sql_caching.logger") @@ -990,7 +1047,6 @@ def test_http_error_logs_s3_diagnostics(self, mock_put, mock_logger): 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",) ) @@ -1075,7 +1131,6 @@ def test_serialization_failure_uses_distinct_cause(self, mock_put, mock_logger): 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") @@ -1095,61 +1150,66 @@ def test_seconds_since_url_issued_reflects_elapsed_time( 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.tempfile.TemporaryFile") + def test_failure_before_the_request_leaves_both_times_unset( + self, mock_temp_file, mock_logger + ): + """Nothing was timed because nothing was sent, and neither may raise.""" + mock_temp_file.side_effect = OSError("no space left on device") + + upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) + + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertIsNone(extra["seconds_since_url_issued"]) + self.assertIsNone(extra["upload_duration_seconds"]) + @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) + @patch("deepnote_toolkit.sql.sql_caching.time.monotonic") + def test_transfer_time_is_not_counted_as_url_age( + self, mock_monotonic, mock_put, mock_logger + ): + clock = [1000.0] + mock_monotonic.side_effect = lambda: clock[0] + + def slow_put(*args, **kwargs): + clock[0] += 1000.0 + return _s3_response(500, b"InternalError") + + mock_put.side_effect = slow_put upload_sql_cache( pd.DataFrame({"a": [1, 2, 3]}), - SqlCacheUpload(url=PRESIGNED_URL, issued_at="not-a-number"), + SqlCacheUpload(url=PRESIGNED_URL, issued_at=100.5), ) extra = mock_logger.error.call_args.kwargs["extra"] - self.assertIsNone(extra["seconds_since_url_issued"]) + # the request left inside the 900s window, so the slow transfer that + # followed must not make the entry read as an expired URL + self.assertEqual(extra["seconds_since_url_issued"], 899.5) + self.assertEqual(extra["upload_duration_seconds"], 1000.0) - @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) + def test_upload_request_bounds_connect_and_read_phases(self, mock_put, mock_logger): + mock_put.return_value = _s3_response(200) + + upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) - if fail_serialization: - dataframe = mock.Mock() - dataframe.to_parquet.side_effect = ValueError("cannot serialize") - else: - dataframe = pd.DataFrame({"a": [1, 2, 3]}) + connect_timeout, read_timeout = mock_put.call_args.kwargs["timeout"] + self.assertIsNotNone(connect_timeout) + self.assertIsNotNone(read_timeout) - self.assertIsNone(upload_sql_cache(dataframe, _upload())) + @patch("deepnote_toolkit.sql.sql_caching.logger") + @patch("deepnote_toolkit.sql.sql_caching.requests.put") + def test_upload_timeout_never_raises(self, mock_put, mock_logger): + mock_put.side_effect = requests.exceptions.Timeout(URLLIB3_ERROR_MESSAGE) + + upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) + + extra = mock_logger.error.call_args.kwargs["extra"] + self.assertEqual(extra["error_type"], "Timeout") class TestLoggedExtras(unittest.TestCase): @@ -1158,12 +1218,17 @@ class TestLoggedExtras(unittest.TestCase): 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_every_failure_path_logs_a_cause(self): + self.assertEqual( + {extra["sql_caching_cause"] for extra in self.extras}, + { + "failed_to_upload_to_cache", + "failed_to_serialize_cache", + "failed_to_download_from_cache", + "failed_to_request_cache_info", + "http_error", + }, + ) def test_extra_keys_avoid_reserved_logrecord_attributes(self): """A reserved key makes logger.error raise into the user's query.""" From 9eb0b3e65f70cc0e159bf6795c7b210270fd1f29 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 19:53:08 +0000 Subject: [PATCH 3/5] refactor: extract SQL cache diagnostics into their own module sql_caching.py had grown a log-redaction subsystem alongside the cache flow. Move the redaction, parsing and formatting helpers to sql_cache_diagnostics.py and colocate their tests, leaving sql_caching.py as the read/write flow it started as. Names crossing the new module boundary drop their leading underscore. The repeated redact-and-bound idiom becomes _redacted_snippet, which keeps the field-length bound private to the new module. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_018cayDdcs1iuui8Pgh5npKm --- deepnote_toolkit/sql/sql_cache_diagnostics.py | 199 +++++++++ deepnote_toolkit/sql/sql_caching.py | 217 +--------- tests/unit/helpers/sql_cache_fixtures.py | 98 +++++ tests/unit/test_sql_cache_diagnostics.py | 265 ++++++++++++ tests/unit/test_sql_caching.py | 385 ++---------------- 5 files changed, 613 insertions(+), 551 deletions(-) create mode 100644 deepnote_toolkit/sql/sql_cache_diagnostics.py create mode 100644 tests/unit/helpers/sql_cache_fixtures.py create mode 100644 tests/unit/test_sql_cache_diagnostics.py diff --git a/deepnote_toolkit/sql/sql_cache_diagnostics.py b/deepnote_toolkit/sql/sql_cache_diagnostics.py new file mode 100644 index 0000000..2e88787 --- /dev/null +++ b/deepnote_toolkit/sql/sql_cache_diagnostics.py @@ -0,0 +1,199 @@ +"""Log-safe descriptions of SQL cache failures. + +Everything here exists so a failed cache request can be explained in a log entry +without that entry carrying the presigned URL's signing parameters with it. +""" + +import re +from typing import Any, Optional +from urllib.parse import parse_qs, urlsplit + +import requests + +_MAX_ERROR_BODY_BYTES = 4096 +_MAX_ERROR_FIELD_CHARS = 500 +_MAX_OBJECT_PATH_CHARS = 200 +_MAX_RAW_EXCEPTION_CHARS = _MAX_ERROR_FIELD_CHARS * 4 + +# Matches the scheme-less, path-only form that urllib3 puts in its connection +# error messages as well as a whole URL +_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, +) +# urlsplit only splits on a literal '?', so an over-encoded URL keeps its signed +# query string in the path +_ENCODED_QUERY_SEPARATOR = re.compile("%3F", re.IGNORECASE) +# 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 SqlCacheHttpError(Exception): + """A non-2xx response from the cache object store. + + Its own string representation deliberately contains no URL. + """ + + def __init__(self, diagnostics: dict[str, Any]) -> None: + """Store log-safe diagnostics already taken from the failed 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.""" + redacted = _URL_QUERY_PATTERN.sub(r"\1?", text) + return _AWS_CREDENTIAL_PARAM_PATTERN.sub(r"\1", redacted) + + +def _redacted_snippet(text: str) -> str: + """Redact text and cut it to what a single log field may carry.""" + return redact_sensitive(text)[:_MAX_ERROR_FIELD_CHARS] + + +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_between(start: Optional[float], end: Optional[float]) -> Optional[float]: + """Monotonic seconds from start to end, or None when either was not taken.""" + if start is None or end is None: + return None + + return round(end - start, 1) + + +def safe_url_path(url: str) -> Optional[str]: + """The bounded path of a URL, or None when the URL cannot be parsed.""" + if not isinstance(url, str): + 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.""" + 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") + path = _ENCODED_QUERY_SEPARATOR.split(parts.path, maxsplit=1)[0] + return { + # hostname rather than netloc, which would carry any user:pass@ userinfo + "object_host": parts.hostname, + # cut for an encoded '?', redacted for separators urlsplit ignores + "object_path": redact_sensitive(path[:_MAX_OBJECT_PATH_CHARS]), + "url_expires_in": _to_int_or_none( + expires_values[0] if expires_values else None + ), + } + except Exception: + 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. + + Taken off the wire rather than sliced off response.content, which would + download the whole body first, and accumulated rather than read once, because + a chunked response answers with one HTTP chunk per read however large a size + is asked for. + """ + try: + chunks = [] + length = 0 + for chunk in response.iter_content(_MAX_ERROR_BODY_BYTES): + chunks.append(chunk) + length += len(chunk) + if length >= _MAX_ERROR_BODY_BYTES: + break + + prefix = b"".join(chunks)[:_MAX_ERROR_BODY_BYTES] + return prefix.decode("utf-8", errors="replace") + except Exception: + return None + + +def read_body_snippet(response: requests.Response) -> Optional[str]: + """A log-safe prefix of a response body, or None when it is unreadable.""" + body = _read_response_body_prefix(response) + return None if body is None else _redacted_snippet(body) + + +def describe_s3_error(response: requests.Response) -> dict[str, Any]: + """Non-sensitive diagnostics from a failed response from the object store. + + Only the allowlisted fields are taken from the body, because the other + elements of an S3 error document (, , + ) echo the signed query string and with it the credentials. + """ + 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] = _redacted_snippet(value) + + if diagnostics["s3_error_code"] is None: + # Not an S3 error document - a proxy or gateway answered instead + diagnostics["response_body_snippet"] = _redacted_snippet(body) + + 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, and an exception message has no bound of its own + return { + "error_type": type(exc).__name__, + "error_message": _redacted_snippet(str(exc)[:_MAX_RAW_EXCEPTION_CHARS]), + } diff --git a/deepnote_toolkit/sql/sql_caching.py b/deepnote_toolkit/sql/sql_caching.py index b23486e..9b9ecc5 100644 --- a/deepnote_toolkit/sql/sql_caching.py +++ b/deepnote_toolkit/sql/sql_caching.py @@ -1,16 +1,23 @@ 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 from pyarrow import ArrowInvalid, ArrowNotImplementedError +from deepnote_toolkit.sql.sql_cache_diagnostics import ( + SqlCacheHttpError, + describe_exception, + describe_presigned_url, + describe_s3_error, + read_body_snippet, + safe_url_path, + seconds_between, +) from deepnote_toolkit.sql.sql_utils import is_single_select_query from ..get_webapp_url import get_absolute_userpod_api_url, get_project_auth_headers @@ -24,31 +31,6 @@ # a whole, so it cannot abort a large but healthy download. _OBJECT_STORE_TIMEOUT: tuple[int, int] = (5, 60) -_MAX_ERROR_BODY_BYTES = 4096 -_MAX_ERROR_FIELD_CHARS = 500 -_MAX_OBJECT_PATH_CHARS = 200 -_MAX_RAW_EXCEPTION_CHARS = _MAX_ERROR_FIELD_CHARS * 4 - -# Matches the scheme-less, path-only form that urllib3 puts in its connection -# error messages as well as a whole URL -_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, -) -# urlsplit only splits on a literal '?', so an over-encoded URL keeps its signed -# query string in the path -_ENCODED_QUERY_SEPARATOR = re.compile("%3F", re.IGNORECASE) -# 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.""" @@ -57,161 +39,6 @@ class SqlCacheUpload(NamedTuple): issued_at: float -class _SqlCacheHttpError(Exception): - """A non-2xx response from the cache object store. - - Its own string representation deliberately contains no URL. - """ - - def __init__(self, diagnostics: dict[str, Any]) -> None: - """Store log-safe diagnostics already taken from the failed 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.""" - 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_between(start: Optional[float], end: Optional[float]) -> Optional[float]: - """Monotonic seconds from start to end, or None when either was not taken.""" - if start is None or end is None: - return None - - return round(end - start, 1) - - -def _safe_url_path(url: str) -> Optional[str]: - """The bounded path of a URL, or None when the URL cannot be parsed.""" - if not isinstance(url, str): - 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.""" - 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") - path = _ENCODED_QUERY_SEPARATOR.split(parts.path, maxsplit=1)[0] - return { - # hostname rather than netloc, which would carry any user:pass@ userinfo - "object_host": parts.hostname, - # cut for an encoded '?', redacted for separators urlsplit ignores - "object_path": _redact_sensitive(path[:_MAX_OBJECT_PATH_CHARS]), - "url_expires_in": _to_int_or_none( - expires_values[0] if expires_values else None - ), - } - except Exception: - 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. - - Taken off the wire rather than sliced off response.content, which would - download the whole body first, and accumulated rather than read once, because - a chunked response answers with one HTTP chunk per read however large a size - is asked for. - """ - try: - chunks = [] - length = 0 - for chunk in response.iter_content(_MAX_ERROR_BODY_BYTES): - chunks.append(chunk) - length += len(chunk) - if length >= _MAX_ERROR_BODY_BYTES: - break - - prefix = b"".join(chunks)[:_MAX_ERROR_BODY_BYTES] - return prefix.decode("utf-8", errors="replace") - except Exception: - return None - - -def _describe_s3_error(response: requests.Response) -> dict[str, Any]: - """Non-sensitive diagnostics from a failed response from the object store. - - Only the allowlisted fields are taken from the body, because the other - elements of an S3 error document (, , - ) echo the signed query string and with it the credentials. - """ - 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 - 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, and an exception message has no bound of its own - 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: str, bind_params: dict, @@ -261,7 +88,7 @@ def get_sql_cache( "Failed to request SQL cache info", extra={ "sql_caching_cause": "failed_to_request_cache_info", - **_describe_exception(exc), + **describe_exception(exc), }, ) return None, None @@ -278,8 +105,8 @@ def get_sql_cache( "Failed to download dataframe from cache", extra={ "sql_caching_cause": "failed_to_download_from_cache", - **_describe_exception(exc), - **_describe_presigned_url(download_url), + **describe_exception(exc), + **describe_presigned_url(download_url), }, ) return None, None @@ -357,22 +184,22 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: # 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) + failure = describe_s3_error(response) except Exception as exc: - failure = _describe_exception(exc) + 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_between( + **describe_presigned_url(upload.url), + "seconds_since_url_issued": seconds_between( upload.issued_at, put_started_at ), # Separate field so "the URL was already old" stays distinguishable # from "the transfer was slow" - "upload_duration_seconds": _seconds_between( + "upload_duration_seconds": seconds_between( put_started_at, time.monotonic() ), }, @@ -391,7 +218,7 @@ def _try_read_cache(download_url: str) -> pd.DataFrame: download_url, timeout=_OBJECT_STORE_TIMEOUT, stream=True ) as response: if response.status_code >= 400: - raise _SqlCacheHttpError(_describe_s3_error(response)) + raise SqlCacheHttpError(describe_s3_error(response)) # Read inside the block: once it exits, the body reads back as empty # rather than raising, which would cache-miss every hit in silence @@ -435,17 +262,13 @@ def _request_cache_info_from_webapp( ) if sql_cache_response.status_code != 200: # the caching endpoint is not available, we can't use it. We'll skip the caching logic - 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, + "cache_info_path": safe_url_path(sql_cache_url), + "response_body_snippet": read_body_snippet(sql_cache_response), }, ) return None diff --git a/tests/unit/helpers/sql_cache_fixtures.py b/tests/unit/helpers/sql_cache_fixtures.py new file mode 100644 index 0000000..d0373a8 --- /dev/null +++ b/tests/unit/helpers/sql_cache_fixtures.py @@ -0,0 +1,98 @@ +"""Stand-ins for what the SQL cache object store answers with. + +Shared by the tests for the cache flow and for the diagnostics taken off its +failures, so both see the same bodies and the same response semantics. +""" + +from unittest import mock + +# Signing 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'))" +) + +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", +} + +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 +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 +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"" +) + + +def s3_response(status_code, body=b"", headers=None): + """Build a stand-in for a requests.Response from the object store.""" + response = mock.MagicMock( + status_code=status_code, content=body, headers=headers or {} + ) + response.__enter__.return_value = response + # a real streamed body reads back empty once the block exits, silently + response.__exit__.side_effect = lambda *_: setattr(response, "content", b"") + # a chunked response yields one HTTP chunk per read however large a size is + # asked for + response.iter_content.side_effect = lambda size: iter( + [body[i : i + 20] for i in range(0, len(body), 20)] + ) + return response diff --git a/tests/unit/test_sql_cache_diagnostics.py b/tests/unit/test_sql_cache_diagnostics.py new file mode 100644 index 0000000..ba80259 --- /dev/null +++ b/tests/unit/test_sql_cache_diagnostics.py @@ -0,0 +1,265 @@ +import json +import time +import unittest +from unittest import mock + +from parameterized import parameterized + +from deepnote_toolkit.sql.sql_cache_diagnostics import ( + _URL_QUERY_PATTERN, + describe_exception, + describe_presigned_url, + describe_s3_error, + redact_sensitive, + safe_url_path, +) + +from .helpers.sql_cache_fixtures import ( + ACCESS_DENIED_EXPIRED_BODY, + AWS_HEADERS, + EXPIRED_TOKEN_BODY, + PRESIGNED_URL, + PROXY_ECHO_BODY, + SECRETS, + SIGNATURE_MISMATCH_BODY, + URLLIB3_ERROR_MESSAGE, + s3_response, +) + + +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") + 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) + + def test_body_prefix_is_streamed_rather_than_buffered(self): + """Reading .content downloads the whole body just to keep 4 KB of it.""" + buffered = [] + response = mock.MagicMock(status_code=403, headers={}) + type(response).content = mock.PropertyMock( + side_effect=lambda: buffered.append("content") + ) + response.iter_content.side_effect = lambda size: iter( + [ACCESS_DENIED_EXPIRED_BODY[:size]] + ) + + diagnostics = describe_s3_error(response) + + self.assertEqual(buffered, []) + self.assertEqual(diagnostics["s3_error_code"], "AccessDenied") + + +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_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"), + ("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)) + + @parameterized.expand( + [ + ( + "separator_encoded", + "%3FX-Amz-Credential=CREDVALUE&X-Amz-Security-Token=TOKENVALUE" + "&X-Amz-Signature=SIGVALUE", + ), + ( + "separator_and_equals_encoded", + "%3FX-Amz-Credential%3DCREDVALUE&X-Amz-Security-Token%3DTOKENVALUE" + "&X-Amz-Signature%3DSIGVALUE", + ), + ( + "whole_query_encoded", + "%3FX-Amz-Credential%3DCREDVALUE%26X-Amz-Security-Token%3DTOKENVALUE" + "%26X-Amz-Signature%3DSIGVALUE", + ), + ("separator_is_a_semicolon", ";X-Amz-Credential=CREDVALUE"), + ("no_separator_at_all", "X-Amz-Signature=SIGVALUE"), + ] + ) + def test_over_encoded_url_does_not_leak_signing_params_via_path(self, _, suffix): + """urlsplit only splits on a literal '?', so the path carries the rest.""" + described = describe_presigned_url( + "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" + suffix + ) + + # the parameter names may survive redaction, their values must not + for value in described.values(): + for secret in ("CREDVALUE", "TOKENVALUE", "SIGVALUE"): + self.assertNotIn(secret, str(value)) diff --git a/tests/unit/test_sql_caching.py b/tests/unit/test_sql_caching.py index 01395aa..05bd2e6 100644 --- a/tests/unit/test_sql_caching.py +++ b/tests/unit/test_sql_caching.py @@ -12,118 +12,33 @@ 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" - -# Signing 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'))" -) - -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", -} - -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" +from .helpers.sql_cache_fixtures import ( + ACCESS_DENIED_EXPIRED_BODY, + AWS_HEADERS, + EXPIRED_TOKEN_BODY, + PRESIGNED_URL, + PROXY_ECHO_BODY, + SECRETS, + SIGNATURE_MISMATCH_BODY, + URLLIB3_ERROR_MESSAGE, + s3_response, ) -# 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 -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 -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"" -) +QUERY = "SELECT * FROM users" 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.""" - response = mock.MagicMock( - status_code=status_code, content=body, headers=headers or {} - ) - response.__enter__.return_value = response - # a real streamed body reads back empty once the block exits, silently - response.__exit__.side_effect = lambda *_: setattr(response, "content", b"") - # a chunked response yields one HTTP chunk per read however large a size is - # asked for - response.iter_content.side_effect = lambda size: iter( - [body[i : i + 20] for i in range(0, len(body), 20)] - ) - return response - - def _upload(url=PRESIGNED_URL, issued_at=None): """Build an upload handle, issued now unless told otherwise.""" return SqlCacheUpload( @@ -162,7 +77,7 @@ def _collect_logged_extras(): with patch("deepnote_toolkit.sql.sql_caching.logger") as mock_logger: with patch("deepnote_toolkit.sql.sql_caching.requests.put") as mock_put: - mock_put.return_value = _s3_response( + mock_put.return_value = s3_response( 403, ACCESS_DENIED_EXPIRED_BODY, AWS_HEADERS ) upload_sql_cache(dataframe, _upload()) @@ -181,7 +96,7 @@ def _collect_logged_extras(): ) as mock_cache_info: mock_cache_info.return_value = _cache_hit() with patch("deepnote_toolkit.sql.sql_caching.requests.get") as mock_get: - mock_get.return_value = _s3_response( + mock_get.return_value = s3_response( 403, EXPIRED_TOKEN_BODY, AWS_HEADERS ) get_sql_cache(QUERY, {}, "123", "read", "dataframe") @@ -206,7 +121,7 @@ def _collect_logged_extras(): "?sqlCacheKey=abc&sqlCacheMode=read" ) mock_headers.return_value = {} - mock_get.return_value = _s3_response(503, b"upstream unavailable") + 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] @@ -347,7 +262,7 @@ 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_get.return_value = s3_response(200, b"parquet-bytes") mock_read_parquet.return_value = pd.DataFrame() result_df, upload = get_sql_cache( @@ -394,7 +309,7 @@ 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_get.return_value = s3_response(200, b"pickle-bytes") mock_read_parquet.side_effect = ArrowInvalid mock_read_pickle.return_value = pd.DataFrame() @@ -446,7 +361,7 @@ 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_get.return_value = s3_response(200, b"parquet-bytes") mock_read_parquet.side_effect = Exception("Failed to download from cache") result_df, upload = get_sql_cache( @@ -545,7 +460,7 @@ def test_read_from_cache_error_doesnt_raise( 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_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") @@ -564,7 +479,7 @@ 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) + mock_get.return_value = s3_response(403, EXPIRED_TOKEN_BODY, AWS_HEADERS) result_df, upload = get_sql_cache(QUERY, {}, "123", "read", "dataframe") @@ -597,7 +512,7 @@ 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) + mock_get.return_value = s3_response(403, body, AWS_HEADERS) get_sql_cache(QUERY, {}, "123", "read", "dataframe") @@ -617,7 +532,7 @@ def test_download_reads_parquet_from_fetched_bytes( dataframe.to_parquet(buffer) mock_cache_info.return_value = _cache_hit() - mock_get.return_value = _s3_response(200, buffer.getvalue()) + mock_get.return_value = s3_response(200, buffer.getvalue()) result_df, _ = get_sql_cache(QUERY, {}, "123", "read", "dataframe") @@ -637,7 +552,7 @@ def test_download_falls_back_to_pickle_from_same_bytes( dataframe.to_pickle(buffer) mock_cache_info.return_value = _cache_hit() - mock_get.return_value = _s3_response(200, buffer.getvalue()) + mock_get.return_value = s3_response(200, buffer.getvalue()) result_df, _ = get_sql_cache(QUERY, {}, "123", "read", "dataframe") @@ -657,7 +572,7 @@ def test_download_request_is_bounded_and_streamed( dataframe.to_parquet(buffer) mock_cache_info.return_value = _cache_hit() - mock_get.return_value = _s3_response(200, buffer.getvalue()) + mock_get.return_value = s3_response(200, buffer.getvalue()) get_sql_cache(QUERY, {}, "123", "read", "dataframe") @@ -703,7 +618,7 @@ def test_non_200_logs_constant_message_with_status_in_extra( "?sqlCacheKey=abc&sqlCacheMode=read" ) mock_headers.return_value = {} - mock_get.return_value = _s3_response(503, b"upstream unavailable") + mock_get.return_value = s3_response(503, b"upstream unavailable") self.assertIsNone(_request_cache_info_from_webapp("abc", "123", "read")) @@ -720,244 +635,6 @@ def test_non_200_logs_constant_message_with_status_in_extra( 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") - 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) - - def test_body_prefix_is_streamed_rather_than_buffered(self): - """Reading .content downloads the whole body just to keep 4 KB of it.""" - buffered = [] - response = mock.MagicMock(status_code=403, headers={}) - type(response).content = mock.PropertyMock( - side_effect=lambda: buffered.append("content") - ) - response.iter_content.side_effect = lambda size: iter( - [ACCESS_DENIED_EXPIRED_BODY[:size]] - ) - - diagnostics = _describe_s3_error(response) - - self.assertEqual(buffered, []) - self.assertEqual(diagnostics["s3_error_code"], "AccessDenied") - - -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_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"), - ("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)) - - @parameterized.expand( - [ - ( - "separator_encoded", - "%3FX-Amz-Credential=CREDVALUE&X-Amz-Security-Token=TOKENVALUE" - "&X-Amz-Signature=SIGVALUE", - ), - ( - "separator_and_equals_encoded", - "%3FX-Amz-Credential%3DCREDVALUE&X-Amz-Security-Token%3DTOKENVALUE" - "&X-Amz-Signature%3DSIGVALUE", - ), - ( - "whole_query_encoded", - "%3FX-Amz-Credential%3DCREDVALUE%26X-Amz-Security-Token%3DTOKENVALUE" - "%26X-Amz-Signature%3DSIGVALUE", - ), - ("separator_is_a_semicolon", ";X-Amz-Credential=CREDVALUE"), - ("no_separator_at_all", "X-Amz-Signature=SIGVALUE"), - ] - ) - def test_over_encoded_url_does_not_leak_signing_params_via_path(self, _, suffix): - """urlsplit only splits on a literal '?', so the path carries the rest.""" - described = _describe_presigned_url( - "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" + suffix - ) - - # the parameter names may survive redaction, their values must not - for value in described.values(): - for secret in ("CREDVALUE", "TOKENVALUE", "SIGVALUE"): - self.assertNotIn(secret, str(value)) - - class TestUploadSqlCache(unittest.TestCase): @patch("deepnote_toolkit.sql.sql_caching.logger") @patch("deepnote_toolkit.sql.sql_caching.requests.put") @@ -1040,7 +717,7 @@ def capture_file_state(f, **_kwargs): @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( + mock_put.return_value = s3_response( 403, ACCESS_DENIED_EXPIRED_BODY, AWS_HEADERS ) @@ -1074,7 +751,7 @@ def test_http_error_logs_s3_diagnostics(self, mock_put, mock_logger): 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) + mock_put.return_value = s3_response(403, body, AWS_HEADERS) upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) @@ -1105,9 +782,9 @@ def test_message_is_constant_across_different_failures(self, mock_put, mock_logg 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) + 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") + 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 @@ -1138,7 +815,7 @@ def test_serialization_failure_uses_distinct_cause(self, mock_put, mock_logger): 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) + mock_put.return_value = s3_response(403, ACCESS_DENIED_EXPIRED_BODY) upload_sql_cache( pd.DataFrame({"a": [1, 2, 3]}), @@ -1175,7 +852,7 @@ def test_transfer_time_is_not_counted_as_url_age( def slow_put(*args, **kwargs): clock[0] += 1000.0 - return _s3_response(500, b"InternalError") + return s3_response(500, b"InternalError") mock_put.side_effect = slow_put @@ -1193,7 +870,7 @@ def slow_put(*args, **kwargs): @patch("deepnote_toolkit.sql.sql_caching.logger") @patch("deepnote_toolkit.sql.sql_caching.requests.put") def test_upload_request_bounds_connect_and_read_phases(self, mock_put, mock_logger): - mock_put.return_value = _s3_response(200) + mock_put.return_value = s3_response(200) upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) From d1f62aa491d8ad1ac71a01b75ee4bc7574b3e96e Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 19:01:34 +0000 Subject: [PATCH 4/5] refactor: address review on the SQL cache diagnostics module Type the URL parameters `object` rather than `str`. Both guard with isinstance because the URLs come out of a JSON response, so annotating them `str` made the guard read as dead code - and mypy does not flag it even with warn_unreachable. Drive the S3 error extraction from one table. The four element names were spelled out three times over, in the regex, the loop and the defaults; adding a field is now one line. Use raise_for_status() on both transfers and describe the failure from the exception. sql_caching no longer interprets status codes and every failure goes through describe_exception, which reads the error document off the response the exception carries. The download response is streamed, so it is described inside the `with` block: once the connection is released the body reads back empty rather than raising, and SqlCacheHttpError carries the finished diagnostics out. Make the response double faithful to that - raise_for_status() raises a real HTTPError carrying the response, and a released body reads back empty through iter_content as well as .content. Without the second part, moving the error handling outside the block passed every test while production logged nothing. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_018cayDdcs1iuui8Pgh5npKm --- deepnote_toolkit/sql/sql_cache_diagnostics.py | 50 +++++++++++-------- deepnote_toolkit/sql/sql_caching.py | 22 ++++---- tests/unit/helpers/sql_cache_fixtures.py | 39 +++++++++++---- 3 files changed, 69 insertions(+), 42 deletions(-) diff --git a/deepnote_toolkit/sql/sql_cache_diagnostics.py b/deepnote_toolkit/sql/sql_cache_diagnostics.py index 2e88787..571291c 100644 --- a/deepnote_toolkit/sql/sql_cache_diagnostics.py +++ b/deepnote_toolkit/sql/sql_cache_diagnostics.py @@ -29,10 +29,21 @@ # urlsplit only splits on a literal '?', so an over-encoded URL keeps its signed # query string in the path _ENCODED_QUERY_SEPARATOR = re.compile("%3F", re.IGNORECASE) +# The only elements ever read out of an S3 error document, mapped to the log +# field each becomes. The rest of the document (, +# , ) echoes the signed query string and with it +# the credentials, so nothing outside this table is touched. # and are AWS's own answer to whether the URL ran out of -# time, independent of this pod's clock +# time, independent of this pod's clock. +_S3_ERROR_FIELDS = { + "Code": "s3_error_code", + "Message": "s3_error_message", + "Expires": "s3_expires", + "ServerTime": "s3_server_time", +} + _S3_ERROR_FIELD_PATTERN = re.compile( - r"<(Code|Message|Expires|ServerTime)>(.*?)", re.DOTALL | re.IGNORECASE + rf"<({'|'.join(_S3_ERROR_FIELDS)})>(.*?)", re.DOTALL | re.IGNORECASE ) @@ -78,7 +89,7 @@ def seconds_between(start: Optional[float], end: Optional[float]) -> Optional[fl return round(end - start, 1) -def safe_url_path(url: str) -> Optional[str]: +def safe_url_path(url: object) -> Optional[str]: """The bounded path of a URL, or None when the URL cannot be parsed.""" if not isinstance(url, str): return None @@ -89,7 +100,7 @@ def safe_url_path(url: str) -> Optional[str]: return None -def describe_presigned_url(url: str) -> dict[str, Any]: +def describe_presigned_url(url: object) -> dict[str, Any]: """Object path and declared expiry, without the credential-bearing query string.""" if not isinstance(url, str): # urlsplit(None) answers with bytes-valued fields instead of raising, and a @@ -145,9 +156,7 @@ def read_body_snippet(response: requests.Response) -> Optional[str]: def describe_s3_error(response: requests.Response) -> dict[str, Any]: """Non-sensitive diagnostics from a failed response from the object store. - Only the allowlisted fields are taken from the body, because the other - elements of an S3 error document (, , - ) echo the signed query string and with it the credentials. + Only the elements in _S3_ERROR_FIELDS are read out of the body. """ diagnostics: dict[str, Any] = { "status_code": response.status_code, @@ -156,28 +165,21 @@ def describe_s3_error(response: requests.Response) -> dict[str, Any]: "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, + **{field: None for field in _S3_ERROR_FIELDS.values()}, } 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) + found = { + element.lower(): value + for element, 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) + for element, field in _S3_ERROR_FIELDS.items(): + value = found.get(element.lower()) if value is not None: - diagnostics[key] = _redacted_snippet(value) + diagnostics[field] = _redacted_snippet(value) if diagnostics["s3_error_code"] is None: # Not an S3 error document - a proxy or gateway answered instead @@ -191,6 +193,12 @@ def describe_exception(exc: BaseException) -> dict[str, Any]: if isinstance(exc, SqlCacheHttpError): return dict(exc.diagnostics) + # raise_for_status() attaches the response, so the error document is still + # there to be read - unlike its own message, which is just the status line + # and the whole presigned URL + if isinstance(exc, requests.HTTPError) and exc.response is not None: + return describe_s3_error(exc.response) + # Cut before redacting rather than after: the URL pattern backtracks over every # start position, and an exception message has no bound of its own return { diff --git a/deepnote_toolkit/sql/sql_caching.py b/deepnote_toolkit/sql/sql_caching.py index 9b9ecc5..e23aaa8 100644 --- a/deepnote_toolkit/sql/sql_caching.py +++ b/deepnote_toolkit/sql/sql_caching.py @@ -13,7 +13,6 @@ SqlCacheHttpError, describe_exception, describe_presigned_url, - describe_s3_error, read_body_snippet, safe_url_path, seconds_between, @@ -179,12 +178,8 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: 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) + response.raise_for_status() + return except Exception as exc: failure = describe_exception(exc) @@ -213,15 +208,18 @@ def _try_read_cache(download_url: str) -> pd.DataFrame: fetches it with urllib and so raises before S3's error body is ever read. """ # Streamed so that a failed download costs a bounded prefix rather than the - # whole error body, and so the rest of it is dropped with the connection + # whole error body, and so the rest of it is dropped with the connection. + # Both branches must read the body inside the block: once it exits the body + # reads back as empty rather than raising, which would silently cache-miss + # every hit and drop S3's error document on the way out. with requests.get( download_url, timeout=_OBJECT_STORE_TIMEOUT, stream=True ) as response: - if response.status_code >= 400: - raise SqlCacheHttpError(describe_s3_error(response)) + try: + response.raise_for_status() + except requests.HTTPError as exc: + raise SqlCacheHttpError(describe_exception(exc)) from exc - # Read inside the block: once it exits, the body reads back as empty - # rather than raising, which would cache-miss every hit in silence buffer = BytesIO(response.content) try: diff --git a/tests/unit/helpers/sql_cache_fixtures.py b/tests/unit/helpers/sql_cache_fixtures.py index d0373a8..c6befe5 100644 --- a/tests/unit/helpers/sql_cache_fixtures.py +++ b/tests/unit/helpers/sql_cache_fixtures.py @@ -6,6 +6,8 @@ from unittest import mock +import requests + # Signing parameters given values that are easy to search log output for PRESIGNED_URL = ( "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" @@ -82,17 +84,36 @@ ) -def s3_response(status_code, body=b"", headers=None): +def s3_response(status_code, body=b"", headers=None, url=PRESIGNED_URL): """Build a stand-in for a requests.Response from the object store.""" response = mock.MagicMock( - status_code=status_code, content=body, headers=headers or {} + status_code=status_code, content=body, headers=headers or {}, url=url ) + released = False + + def release(*_): + nonlocal released + released = True + response.content = b"" + + def iter_content(size): + # a chunked response yields one HTTP chunk per read however large a size + # is asked for + if released: + return iter([]) + return iter([body[i : i + 20] for i in range(0, len(body), 20)]) + response.__enter__.return_value = response - # a real streamed body reads back empty once the block exits, silently - response.__exit__.side_effect = lambda *_: setattr(response, "content", b"") - # a chunked response yields one HTTP chunk per read however large a size is - # asked for - response.iter_content.side_effect = lambda size: iter( - [body[i : i + 20] for i in range(0, len(body), 20)] - ) + # a real streamed body reads back empty through either route once the block + # exits and the connection goes back to the pool - empty, not raising + response.__exit__.side_effect = release + response.iter_content.side_effect = iter_content + if status_code >= 400: + # the real message interpolates the whole requested URL, query string + # included, and carries the response so the body survives the raise + response.raise_for_status.side_effect = requests.HTTPError( + f"{status_code} Client Error: for url: {url}", response=response + ) + else: + response.raise_for_status.return_value = None return response From c424f02b74dceeb28756857227403cb144d21d34 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 19:17:20 +0000 Subject: [PATCH 5/5] refactor: tighten SQL cache comments and upload error logging Trim narrating docstrings and comments while keeping security and regression rationale; log upload failures only in the except path. Co-authored-by: Cursor --- deepnote_toolkit/sql/sql_cache_diagnostics.py | 69 ++++----------- deepnote_toolkit/sql/sql_caching.py | 87 +++++-------------- tests/unit/helpers/sql_cache_fixtures.py | 25 ++---- tests/unit/test_sql_cache_diagnostics.py | 29 ++----- tests/unit/test_sql_caching.py | 30 +------ 5 files changed, 55 insertions(+), 185 deletions(-) diff --git a/deepnote_toolkit/sql/sql_cache_diagnostics.py b/deepnote_toolkit/sql/sql_cache_diagnostics.py index 571291c..62f81b7 100644 --- a/deepnote_toolkit/sql/sql_cache_diagnostics.py +++ b/deepnote_toolkit/sql/sql_cache_diagnostics.py @@ -1,7 +1,6 @@ """Log-safe descriptions of SQL cache failures. -Everything here exists so a failed cache request can be explained in a log entry -without that entry carrying the presigned URL's signing parameters with it. +Failed cache requests are described without presigned URL signing parameters. """ import re @@ -15,26 +14,18 @@ _MAX_OBJECT_PATH_CHARS = 200 _MAX_RAW_EXCEPTION_CHARS = _MAX_ERROR_FIELD_CHARS * 4 -# Matches the scheme-less, path-only form that urllib3 puts in its connection -# error messages as well as a whole URL +# urllib3 connection errors use path-only URLs as well as full URLs _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 +# SignatureDoesNotMatch bodies echo params outside any URL prefix _AWS_CREDENTIAL_PARAM_PATTERN = re.compile( r"(X-(?:Amz|Goog)-(?:Credential|Security-Token|Signature)=)[^&\s\"'<>]*", re.IGNORECASE, ) -# urlsplit only splits on a literal '?', so an over-encoded URL keeps its signed -# query string in the path +# urlsplit only splits on a literal '?' _ENCODED_QUERY_SEPARATOR = re.compile("%3F", re.IGNORECASE) -# The only elements ever read out of an S3 error document, mapped to the log -# field each becomes. The rest of the document (, -# , ) echoes the signed query string and with it -# the credentials, so nothing outside this table is touched. -# and are AWS's own answer to whether the URL ran out of -# time, independent of this pod's clock. +# Parsing stops at this allowlist; other S3 error elements embed the signed query. _S3_ERROR_FIELDS = { "Code": "s3_error_code", "Message": "s3_error_message", @@ -48,30 +39,24 @@ class SqlCacheHttpError(Exception): - """A non-2xx response from the cache object store. - - Its own string representation deliberately contains no URL. - """ + """Non-2xx from the cache object store; str() carries no URL.""" def __init__(self, diagnostics: dict[str, Any]) -> None: - """Store log-safe diagnostics already taken from the failed 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.""" + """Remove credential-bearing material from text destined for logs.""" redacted = _URL_QUERY_PATTERN.sub(r"\1?", text) return _AWS_CREDENTIAL_PARAM_PATTERN.sub(r"\1", redacted) def _redacted_snippet(text: str) -> str: - """Redact text and cut it to what a single log field may carry.""" return redact_sensitive(text)[:_MAX_ERROR_FIELD_CHARS] 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 @@ -82,7 +67,7 @@ def _to_int_or_none(value: Optional[str]) -> Optional[int]: def seconds_between(start: Optional[float], end: Optional[float]) -> Optional[float]: - """Monotonic seconds from start to end, or None when either was not taken.""" + """Elapsed monotonic seconds, or None if either timestamp is missing.""" if start is None or end is None: return None @@ -90,7 +75,7 @@ def seconds_between(start: Optional[float], end: Optional[float]) -> Optional[fl def safe_url_path(url: object) -> Optional[str]: - """The bounded path of a URL, or None when the URL cannot be parsed.""" + """Bounded URL path, or None when parsing fails.""" if not isinstance(url, str): return None @@ -101,10 +86,9 @@ def safe_url_path(url: object) -> Optional[str]: def describe_presigned_url(url: object) -> dict[str, Any]: - """Object path and declared expiry, without the credential-bearing query string.""" + """Object path and declared expiry without the signing query string.""" 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 + # urlsplit(None) returns bytes fields; bytes in logging extras drop the report return {"object_host": None, "object_path": None, "url_expires_in": None} try: @@ -112,9 +96,8 @@ def describe_presigned_url(url: object) -> dict[str, Any]: expires_values = parse_qs(parts.query).get("X-Amz-Expires") path = _ENCODED_QUERY_SEPARATOR.split(parts.path, maxsplit=1)[0] return { - # hostname rather than netloc, which would carry any user:pass@ userinfo + # netloc would include user:pass@ userinfo "object_host": parts.hostname, - # cut for an encoded '?', redacted for separators urlsplit ignores "object_path": redact_sensitive(path[:_MAX_OBJECT_PATH_CHARS]), "url_expires_in": _to_int_or_none( expires_values[0] if expires_values else None @@ -125,16 +108,10 @@ def describe_presigned_url(url: object) -> dict[str, Any]: def _read_response_body_prefix(response: requests.Response) -> Optional[str]: - """Decode a bounded prefix of a response body, or None when it is unreadable. - - Taken off the wire rather than sliced off response.content, which would - download the whole body first, and accumulated rather than read once, because - a chunked response answers with one HTTP chunk per read however large a size - is asked for. - """ try: chunks = [] length = 0 + # iter_content bounds wire read; .content would download the full body for chunk in response.iter_content(_MAX_ERROR_BODY_BYTES): chunks.append(chunk) length += len(chunk) @@ -148,22 +125,17 @@ def _read_response_body_prefix(response: requests.Response) -> Optional[str]: def read_body_snippet(response: requests.Response) -> Optional[str]: - """A log-safe prefix of a response body, or None when it is unreadable.""" + """Redacted response body prefix, or None when unreadable.""" body = _read_response_body_prefix(response) return None if body is None else _redacted_snippet(body) def describe_s3_error(response: requests.Response) -> dict[str, Any]: - """Non-sensitive diagnostics from a failed response from the object store. - - Only the elements in _S3_ERROR_FIELDS are read out of the body. - """ + """Log-safe fields from a failed object-store HTTP response.""" 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"), **{field: None for field in _S3_ERROR_FIELDS.values()}, } @@ -182,25 +154,22 @@ def describe_s3_error(response: requests.Response) -> dict[str, Any]: diagnostics[field] = _redacted_snippet(value) if diagnostics["s3_error_code"] is None: - # Not an S3 error document - a proxy or gateway answered instead + # Non-S3 body (proxy/gateway) diagnostics["response_body_snippet"] = _redacted_snippet(body) return diagnostics def describe_exception(exc: BaseException) -> dict[str, Any]: - """Non-sensitive diagnostics from an exception raised while using the cache.""" + """Log-safe fields from a cache-related exception.""" if isinstance(exc, SqlCacheHttpError): return dict(exc.diagnostics) - # raise_for_status() attaches the response, so the error document is still - # there to be read - unlike its own message, which is just the status line - # and the whole presigned URL + # HTTPError.message is status + URL; the response body is still available if isinstance(exc, requests.HTTPError) and exc.response is not None: return describe_s3_error(exc.response) - # Cut before redacting rather than after: the URL pattern backtracks over every - # start position, and an exception message has no bound of its own + # Truncate before redact: _URL_QUERY_PATTERN is O(n²) on long unbounded text return { "error_type": type(exc).__name__, "error_message": _redacted_snippet(str(exc)[:_MAX_RAW_EXCEPTION_CHARS]), diff --git a/deepnote_toolkit/sql/sql_caching.py b/deepnote_toolkit/sql/sql_caching.py index e23aaa8..c69de4e 100644 --- a/deepnote_toolkit/sql/sql_caching.py +++ b/deepnote_toolkit/sql/sql_caching.py @@ -23,7 +23,6 @@ from ..ipython_utils import output_sql_metadata from ..logging import get_logger -# Initialize logger logger = get_logger() # The read timeout bounds the gap between two received chunks, not the transfer as @@ -45,29 +44,10 @@ def get_sql_cache( sql_cache_mode: str, return_variable_type: str, ) -> tuple[Optional[pd.DataFrame], Optional[SqlCacheUpload]]: - """ - Retrieves the SQL cache from webapp for a given query. - - Args: - query (str): The SQL query to retrieve the cache for. - 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 - pending upload (if applicable). - """ + """Return a cache hit dataframe and/or a presigned upload for a cache miss.""" if not is_single_select_query(query): - # we only cache single select queries - output_sql_metadata( - { - "status": "cache_not_supported_for_query", - # We don't include the additional metadata as the query hasn't been executed/read from cache - } - ) + output_sql_metadata({"status": "cache_not_supported_for_query"}) return None, None query_hash = _generate_cache_key(query, bind_params) @@ -82,7 +62,6 @@ def get_sql_cache( query_hash, integration_id, sql_cache_mode ) except Exception as exc: - # we failed to request the cache info from the webapp logger.error( "Failed to request SQL cache info", extra={ @@ -99,7 +78,6 @@ def get_sql_cache( try: dataframe_from_cache = _try_read_cache(download_url) except Exception as exc: - # we failed to download the dataframe from the cache logger.error( "Failed to download dataframe from cache", extra={ @@ -137,21 +115,14 @@ def _serialize_dataframe_for_cache( 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 + # NB-1684: pickle when parquet cannot represent the frame (e.g. int64 overflow). 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. - """ + """Best-effort cache upload; failures are logged and never raised.""" put_started_at: Optional[float] = None try: @@ -159,8 +130,7 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: try: _serialize_dataframe_for_cache(dataframe, temp_file) except Exception as exc: - # Only the type is logged: serialization errors quote the user's - # column names + # Serialization errors embed column names; log type only. logger.error( "Failed to upload SQL cache", extra={ @@ -171,34 +141,28 @@ def upload_sql_cache(dataframe: pd.DataFrame, upload: SqlCacheUpload) -> None: return temp_file.seek(0) - # S3 checks a presigned signature when the request arrives rather than - # when it finishes, so this is the age that decided whether it was valid + # Presigned validity is checked when the PUT starts, not when it ends. put_started_at = time.monotonic() response = requests.put( upload.url, data=temp_file, timeout=_OBJECT_STORE_TIMEOUT ) response.raise_for_status() - return except Exception as exc: - 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_between( - upload.issued_at, put_started_at - ), - # Separate field so "the URL was already old" stays distinguishable - # from "the transfer was slow" - "upload_duration_seconds": seconds_between( - put_started_at, time.monotonic() - ), - }, - ) + logger.error( + "Failed to upload SQL cache", + extra={ + "sql_caching_cause": "failed_to_upload_to_cache", + **describe_exception(exc), + **describe_presigned_url(upload.url), + "seconds_since_url_issued": seconds_between( + upload.issued_at, put_started_at + ), + "upload_duration_seconds": seconds_between( + put_started_at, time.monotonic() + ), + }, + ) def _try_read_cache(download_url: str) -> pd.DataFrame: @@ -207,11 +171,7 @@ def _try_read_cache(download_url: str) -> pd.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. """ - # Streamed so that a failed download costs a bounded prefix rather than the - # whole error body, and so the rest of it is dropped with the connection. - # Both branches must read the body inside the block: once it exits the body - # reads back as empty rather than raising, which would silently cache-miss - # every hit and drop S3's error document on the way out. + # Read the body before leaving the context: after close, content reads empty. with requests.get( download_url, timeout=_OBJECT_STORE_TIMEOUT, stream=True ) as response: @@ -223,13 +183,8 @@ def _try_read_cache(download_url: str) -> pd.DataFrame: buffer = BytesIO(response.content) try: - # Attempt to read as a parquet file 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 - # pickle format and we don't know which one it is, the file has no extension. - # (see .to_pickle fallback in upload_sql_cache) pass buffer.seek(0) diff --git a/tests/unit/helpers/sql_cache_fixtures.py b/tests/unit/helpers/sql_cache_fixtures.py index c6befe5..91419a5 100644 --- a/tests/unit/helpers/sql_cache_fixtures.py +++ b/tests/unit/helpers/sql_cache_fixtures.py @@ -1,14 +1,8 @@ -"""Stand-ins for what the SQL cache object store answers with. - -Shared by the tests for the cache flow and for the diagnostics taken off its -failures, so both see the same bodies and the same response semantics. -""" - from unittest import mock import requests -# Signing parameters given values that are easy to search log output for +# Distinct tokens so leak assertions can grep log output. PRESIGNED_URL = ( "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" "?X-Amz-Algorithm=AWS4-HMAC-SHA256" @@ -17,8 +11,7 @@ "&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 embeds a path-only URL in the message (no scheme/host). URLLIB3_ERROR_MESSAGE = ( "HTTPSConnectionPool(host='bucket.s3.eu-west-1.amazonaws.com', port=443): " "Max retries exceeded with url: /ws/int/key" @@ -46,7 +39,6 @@ b"REQ123HOSTID456" ) -# S3 rejecting credentials that died before the URL's own expiry EXPIRED_TOKEN_BODY = ( b'\n' b"ExpiredToken" @@ -54,8 +46,7 @@ b"REQ123HOSTID456" ) -# S3 echoes the canonical request - and with it the signed query string - when -# the signature does not match +# Real S3 bodies can embed the signed query in . SIGNATURE_MISMATCH_BODY = ( b'\n' b"SignatureDoesNotMatch" @@ -70,7 +61,7 @@ b"REQ123HOSTID456" ) -# A gateway that answered instead of S3 and echoed the request line back +# Non-S3 gateways may echo the full presigned URL in HTML. PROXY_ECHO_BODY = ( b"502 Bad Gateway\n" b"

502 Bad Gateway

\n" @@ -85,7 +76,6 @@ def s3_response(status_code, body=b"", headers=None, url=PRESIGNED_URL): - """Build a stand-in for a requests.Response from the object store.""" response = mock.MagicMock( status_code=status_code, content=body, headers=headers or {}, url=url ) @@ -97,20 +87,15 @@ def release(*_): response.content = b"" def iter_content(size): - # a chunked response yields one HTTP chunk per read however large a size - # is asked for if released: return iter([]) return iter([body[i : i + 20] for i in range(0, len(body), 20)]) response.__enter__.return_value = response - # a real streamed body reads back empty through either route once the block - # exits and the connection goes back to the pool - empty, not raising + # After context exit, requests clears streamed body content. response.__exit__.side_effect = release response.iter_content.side_effect = iter_content if status_code >= 400: - # the real message interpolates the whole requested URL, query string - # included, and carries the response so the body survives the raise response.raise_for_status.side_effect = requests.HTTPError( f"{status_code} Client Error: for url: {url}", response=response ) diff --git a/tests/unit/test_sql_cache_diagnostics.py b/tests/unit/test_sql_cache_diagnostics.py index ba80259..203fe1d 100644 --- a/tests/unit/test_sql_cache_diagnostics.py +++ b/tests/unit/test_sql_cache_diagnostics.py @@ -29,7 +29,6 @@ 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) @@ -38,7 +37,7 @@ def test_strips_query_string_from_urlopen_style_message(self): 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.""" + """URL query stripping must redact without the AWS-param backstop.""" stripped = _URL_QUERY_PATTERN.sub(r"\1?", URLLIB3_ERROR_MESSAGE) for secret in SECRETS: @@ -69,12 +68,7 @@ def test_leaves_ordinary_text_unchanged(self, _, 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. - """ + """Truncate before redact: _URL_QUERY_PATTERN is quadratic on long strings.""" message = URLLIB3_ERROR_MESSAGE + "&padding=" + "A" * 64_000 started = time.monotonic() @@ -83,8 +77,6 @@ def test_large_message_is_bounded_and_redacted_in_bounded_time(self): 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"]) @@ -100,7 +92,6 @@ def test_extracts_code_and_message_from_xml(self): 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") @@ -122,11 +113,7 @@ def test_captures_aws_request_headers(self): 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. - """ + """Field allowlist must ignore (signed query lives there).""" diagnostics = describe_s3_error( s3_response(403, SIGNATURE_MISMATCH_BODY, AWS_HEADERS) ) @@ -138,12 +125,11 @@ def test_signature_mismatch_body_surfaces_only_code_and_message(self): 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.""" + """Non-S3 bodies are not on the XML allowlist; snippet redaction must apply.""" 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: @@ -176,7 +162,7 @@ def test_oversized_body_is_bounded(self): self.assertLessEqual(len(value), 500) def test_body_prefix_is_streamed_rather_than_buffered(self): - """Reading .content downloads the whole body just to keep 4 KB of it.""" + """Must not touch Response.content (downloads the full body).""" buffered = [] response = mock.MagicMock(status_code=403, headers={}) type(response).content = mock.PropertyMock( @@ -228,7 +214,7 @@ def test_malformed_url_does_not_raise_or_leak(self, _, url): 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 + # Log extras must json.dumps; bytes values would drop the whole report. json.dumps(described) json.dumps(safe_url_path(url)) @@ -254,12 +240,11 @@ def test_malformed_url_does_not_raise_or_leak(self, _, url): ] ) def test_over_encoded_url_does_not_leak_signing_params_via_path(self, _, suffix): - """urlsplit only splits on a literal '?', so the path carries the rest.""" + """urlsplit only splits on a literal '?'; signing params can land in path.""" described = describe_presigned_url( "https://bucket.s3.eu-west-1.amazonaws.com/ws/int/key" + suffix ) - # the parameter names may survive redaction, their values must not for value in described.values(): for secret in ("CREDVALUE", "TOKENVALUE", "SIGVALUE"): self.assertNotIn(secret, str(value)) diff --git a/tests/unit/test_sql_caching.py b/tests/unit/test_sql_caching.py index 05bd2e6..aecd82f 100644 --- a/tests/unit/test_sql_caching.py +++ b/tests/unit/test_sql_caching.py @@ -40,14 +40,12 @@ 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, @@ -56,7 +54,6 @@ def _cache_hit(download_url=PRESIGNED_URL): 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) @@ -67,11 +64,7 @@ def _logged_strings(mock_logger): 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. - """ + """Exercise every error path once for shared extra-dict guards.""" dataframe = pd.DataFrame({"a": [1, 2, 3]}) connection_error = requests.exceptions.ConnectionError(URLLIB3_ERROR_MESSAGE) @@ -131,7 +124,6 @@ class TestGenerateCacheKey(unittest.TestCase): def test_empty_params_returns_valid_result(self): result = _generate_cache_key("SELECT * FROM users", {}) - # assert that the result contains only alphanumeric characters self.assertTrue(result.isalnum()) def test_different_order_of_params_produces_same_result(self): @@ -466,7 +458,6 @@ def test_read_from_cache_error_doesnt_raise( result_df, upload = get_sql_cache(QUERY, {}, "123", "read", "dataframe") - # 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) @@ -499,9 +490,7 @@ def test_download_http_error_logs_s3_diagnostics( @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), ] ) @@ -537,7 +526,6 @@ def test_download_reads_parquet_from_fetched_bytes( 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) @@ -556,9 +544,7 @@ def test_download_falls_back_to_pickle_from_same_bytes( 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.output_sql_metadata") @@ -576,11 +562,9 @@ def test_download_request_is_bounded_and_streamed( get_sql_cache(QUERY, {}, "123", "read", "dataframe") - # an unbounded read phase hangs the user's cell on a socket gone silent connect_timeout, read_timeout = mock_get.call_args.kwargs["timeout"] self.assertIsNotNone(connect_timeout) self.assertIsNotNone(read_timeout) - # and an unstreamed error body costs the size of that body self.assertTrue(mock_get.call_args.kwargs["stream"]) @patch("deepnote_toolkit.sql.sql_caching.logger") @@ -740,9 +724,7 @@ def test_http_error_logs_s3_diagnostics(self, mock_put, mock_logger): @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), ] ) @@ -823,7 +805,6 @@ def test_seconds_since_url_issued_reflects_elapsed_time( ) 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) @@ -832,7 +813,6 @@ def test_seconds_since_url_issued_reflects_elapsed_time( def test_failure_before_the_request_leaves_both_times_unset( self, mock_temp_file, mock_logger ): - """Nothing was timed because nothing was sent, and neither may raise.""" mock_temp_file.side_effect = OSError("no space left on device") upload_sql_cache(pd.DataFrame({"a": [1, 2, 3]}), _upload()) @@ -862,8 +842,6 @@ def slow_put(*args, **kwargs): ) extra = mock_logger.error.call_args.kwargs["extra"] - # the request left inside the 900s window, so the slow transfer that - # followed must not make the entry read as an expired URL self.assertEqual(extra["seconds_since_url_issued"], 899.5) self.assertEqual(extra["upload_duration_seconds"], 1000.0) @@ -890,8 +868,6 @@ def test_upload_timeout_never_raises(self, mock_put, mock_logger): class TestLoggedExtras(unittest.TestCase): - """Guards that apply to every extra dict this module logs.""" - def setUp(self): self.extras = _collect_logged_extras() @@ -908,12 +884,12 @@ def test_every_failure_path_logs_a_cause(self): ) def test_extra_keys_avoid_reserved_logrecord_attributes(self): - """A reserved key makes logger.error raise into the user's query.""" + """Reserved extra keys make logger.error raise into the user's cell.""" 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.""" + """Non-JSON-serializable extras drop the whole error report.""" for extra in self.extras: for value in extra.values(): self.assertIsInstance(value, (str, int, float, bool, type(None)))