feat(sql_execution): Implement retry logic for userpod API#89
feat(sql_execution): Implement retry logic for userpod API#89
Conversation
- Added a new function `_create_retry_session` to create a requests session with retry capabilities for handling 5xx errors on POST requests. - Updated `_generate_temporary_credentials` and `_get_federated_auth_credentials` to use the new retry session for making requests. - Introduced unit tests to verify the retry session configuration and its usage in the credential generation functions.
📝 WalkthroughWalkthroughAdded Sequence Diagram(s)mermaid Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
|
📦 Python package built successfully!
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepnote_toolkit/sql/sql_execution.py`:
- Around line 267-278: The cache lookup in _request_cache_info_from_webapp is
using requests.get directly and thus misses the retry behavior defined in
_create_retry_session; modify _create_retry_session (or its Retry config) to
include GET in allowed_methods (or use allowed_methods=None to cover all
idempotent methods) and then replace the direct requests.get call inside
_request_cache_info_from_webapp with a call through a session returned by
_create_retry_session (e.g., session.get(...)) so cache lookups get the same 5xx
retry protection as POSTs.
- Line 281: Update the signature of
_generate_temporary_credentials(integration_id) to include an explicit return
type annotation (e.g., -> Dict[str, Any]) and add the corresponding typing
imports (from typing import Dict, Any) at the top of the module; ensure the
annotated type matches the actual returned structure from
_generate_temporary_credentials so the function signature and implementation
remain consistent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bf6b45a7-61ec-4a87-a583-c7d418b8fe72
📒 Files selected for processing (2)
deepnote_toolkit/sql/sql_execution.pytests/unit/test_sql_execution.py
| def _create_retry_session() -> requests.Session: | ||
| """Create a requests session with retry on 5xx for POST requests.""" | ||
| session = requests.Session() | ||
| retries = Retry( | ||
| total=3, | ||
| backoff_factor=0.5, | ||
| status_forcelist=[500, 502, 503, 504], | ||
| allowed_methods=["POST"], | ||
| ) | ||
| session.mount("http://", HTTPAdapter(max_retries=retries)) | ||
| session.mount("https://", HTTPAdapter(max_retries=retries)) | ||
| return session |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider extending retry coverage to GET requests in sql_caching.py.
_request_cache_info_from_webapp (line 156 in sql_caching.py) uses requests.get() directly without retry protection. Cache lookups remain vulnerable to transient 5xx errors that POST requests now handle gracefully.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepnote_toolkit/sql/sql_execution.py` around lines 267 - 278, The cache
lookup in _request_cache_info_from_webapp is using requests.get directly and
thus misses the retry behavior defined in _create_retry_session; modify
_create_retry_session (or its Retry config) to include GET in allowed_methods
(or use allowed_methods=None to cover all idempotent methods) and then replace
the direct requests.get call inside _request_cache_info_from_webapp with a call
through a session returned by _create_retry_session (e.g., session.get(...)) so
cache lookups get the same 5xx retry protection as POSTs.
| return session | ||
|
|
||
|
|
||
| def _generate_temporary_credentials(integration_id): |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Add return type annotation.
Missing type hint per coding guidelines.
Proposed fix
-def _generate_temporary_credentials(integration_id):
+def _generate_temporary_credentials(integration_id: str) -> tuple[str, str]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _generate_temporary_credentials(integration_id): | |
| def _generate_temporary_credentials(integration_id: str) -> tuple[str, str]: |
🧰 Tools
🪛 Ruff (0.15.7)
[warning] 281-281: Missing return type annotation for private function _generate_temporary_credentials
(ANN202)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepnote_toolkit/sql/sql_execution.py` at line 281, Update the signature of
_generate_temporary_credentials(integration_id) to include an explicit return
type annotation (e.g., -> Dict[str, Any]) and add the corresponding typing
imports (from typing import Dict, Any) at the top of the module; ensure the
annotated type matches the actual returned structure from
_generate_temporary_credentials so the function signature and implementation
remain consistent.
❌ 1 Tests Failed:
View the top 1 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
|
🚀 Review App Deployment Started
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/unit/test_sql_execution.py`:
- Around line 1055-1065: Convert the context-manager style mocking of
deepnote_toolkit.sql.sql_execution._create_retry_session (currently using with
mock.patch(...) as mock_create_session in the test block) to the decorator style
`@mock.patch` on the test function signature (e.g., add
`@mock.patch`("deepnote_toolkit.sql.sql_execution._create_retry_session") and a
mock_create_session parameter) and remove the inner with-block; apply the same
change to the similar block around lines 1089-1099 so both tests use
decorator-style mocking for consistency with
test_get_federated_auth_credentials_returns_validated_response.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 469a19f4-668a-4b55-8dfc-c8551fc61827
📒 Files selected for processing (1)
tests/unit/test_sql_execution.py
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
deepnote_toolkit/sql/sql_execution.py (1)
280-280:⚠️ Potential issue | 🟡 MinorAdd parameter type hint.
Return type added, but
integration_idlacks type annotation per coding guidelines.🔧 Fix
-def _generate_temporary_credentials(integration_id) -> tuple[str, str]: +def _generate_temporary_credentials(integration_id: str) -> tuple[str, str]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepnote_toolkit/sql/sql_execution.py` at line 280, The function _generate_temporary_credentials has a return type but the parameter integration_id is missing a type annotation; update the function signature for _generate_temporary_credentials to add the appropriate type (e.g., integration_id: str) to comply with the coding guidelines and any static type checks, keeping the return annotation tuple[str, str] unchanged and adjusting any callers only if they rely on a different type.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepnote_toolkit/sql/sql_execution.py`:
- Around line 270-274: The Retry instance named "retries" is missing
allowed_methods so POST requests won't be retried; update the Retry(...) call
(the retries variable in sql_execution.py) to include an allowed_methods
parameter that includes "POST" (for example
allowed_methods=frozenset({"GET","POST","PUT","DELETE","HEAD","OPTIONS","TRACE"})
or at minimum allowed_methods=frozenset({"POST"}) depending on desired scope) so
POST requests are retried per the backoff/status_forcelist settings.
In `@tests/unit/test_sql_execution.py`:
- Around line 1026-1040: The Retry configuration in _create_retry_session is
missing the allowed_methods parameter so POST retries aren't enabled; update the
Retry(...) construction inside the _create_retry_session function to include
allowed_methods=["DELETE","GET","HEAD","OPTIONS","POST","PUT","TRACE"] (or an
equivalent set that includes "POST"), then ensure that same Retry instance is
used when mounting the HTTP/HTTPS adapters so
adapter.max_retries.allowed_methods contains "POST".
---
Duplicate comments:
In `@deepnote_toolkit/sql/sql_execution.py`:
- Line 280: The function _generate_temporary_credentials has a return type but
the parameter integration_id is missing a type annotation; update the function
signature for _generate_temporary_credentials to add the appropriate type (e.g.,
integration_id: str) to comply with the coding guidelines and any static type
checks, keeping the return annotation tuple[str, str] unchanged and adjusting
any callers only if they rely on a different type.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: db1a3ece-0680-4ebe-a62e-d27e219521b2
📒 Files selected for processing (2)
deepnote_toolkit/sql/sql_execution.pytests/unit/test_sql_execution.py
| retries = Retry( | ||
| total=3, | ||
| backoff_factor=0.5, | ||
| status_forcelist=[500, 502, 503, 504], | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
What are the default allowed_methods in urllib3 Retry class and does it include POST?
💡 Result:
The default allowed_methods in the urllib3 Retry class are the idempotent HTTP methods: DELETE, GET, HEAD, OPTIONS, PUT, TRACE. It does not include POST, as POST is not considered idempotent by default.
Citations:
- 1: https://urllib3.readthedocs.io/en/stable/reference/urllib3.util.html
- 2: https://urllib3.readthedocs.io/en/2.3.0/reference/urllib3.util.html
- 3: https://github.com/urllib3/urllib3/blob/main/src/urllib3/util/retry.py
- 4: https://urllib3.readthedocs.io/en/1.26.17/reference/urllib3.util.html
🏁 Script executed:
# Check the actual code at the mentioned lines
cat -n deepnote_toolkit/sql/sql_execution.py | sed -n '260,280p'Repository: deepnote/deepnote-toolkit
Length of output: 878
🏁 Script executed:
# Check the test file for the assertion mentioned
cat -n deepnote_toolkit/tests/test_sql_execution.py | sed -n '1035,1045p'Repository: deepnote/deepnote-toolkit
Length of output: 145
🏁 Script executed:
# Find test files related to sql_execution
find . -name "*test*sql*" -o -name "*sql*test*" | head -20Repository: deepnote/deepnote-toolkit
Length of output: 422
🏁 Script executed:
# Also check the structure of test directories
find . -type d -name "test*" -o -type d -name "*test" | head -10Repository: deepnote/deepnote-toolkit
Length of output: 76
🏁 Script executed:
# Check the test file around line 1040
wc -l tests/unit/test_sql_execution.pyRepository: deepnote/deepnote-toolkit
Length of output: 106
🏁 Script executed:
# Search for allowed_methods in test file
rg -n "allowed_methods" tests/unit/test_sql_execution.pyRepository: deepnote/deepnote-toolkit
Length of output: 132
🏁 Script executed:
# Get context around line 1040 in the test file
sed -n '1030,1050p' tests/unit/test_sql_execution.pyRepository: deepnote/deepnote-toolkit
Length of output: 1120
POST requests won't be retried without explicit allowed_methods.
urllib3's Retry defaults exclude POST from allowed_methods (only includes DELETE, GET, HEAD, OPTIONS, PUT, TRACE). The current config will only retry those methods, not POST despite the docstring stating intent to retry POST.
🐛 Fix: Add allowed_methods
retries = Retry(
total=3,
backoff_factor=0.5,
status_forcelist=[500, 502, 503, 504],
+ allowed_methods=["POST"],
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| retries = Retry( | |
| total=3, | |
| backoff_factor=0.5, | |
| status_forcelist=[500, 502, 503, 504], | |
| ) | |
| retries = Retry( | |
| total=3, | |
| backoff_factor=0.5, | |
| status_forcelist=[500, 502, 503, 504], | |
| allowed_methods=["POST"], | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepnote_toolkit/sql/sql_execution.py` around lines 270 - 274, The Retry
instance named "retries" is missing allowed_methods so POST requests won't be
retried; update the Retry(...) call (the retries variable in sql_execution.py)
to include an allowed_methods parameter that includes "POST" (for example
allowed_methods=frozenset({"GET","POST","PUT","DELETE","HEAD","OPTIONS","TRACE"})
or at minimum allowed_methods=frozenset({"POST"}) depending on desired scope) so
POST requests are retried per the backoff/status_forcelist settings.
| class TestCreateRetrySession(unittest.TestCase): | ||
| def test_retry_session_has_correct_config(self): | ||
| """Test that _create_retry_session configures retries correctly.""" | ||
| from deepnote_toolkit.sql.sql_execution import _create_retry_session | ||
|
|
||
| session = _create_retry_session() | ||
|
|
||
| # Check that both http and https adapters are mounted with retry config | ||
| for prefix in ("http://", "https://"): | ||
| adapter = session.get_adapter(prefix) | ||
| retries = adapter.max_retries | ||
| self.assertEqual(retries.total, 3) | ||
| self.assertEqual(retries.backoff_factor, 0.5) | ||
| self.assertEqual(list(retries.status_forcelist), [500, 502, 503, 504]) | ||
| self.assertIn("POST", retries.allowed_methods) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Run the specific test to confirm failure
cd "$(git rev-parse --show-toplevel)"
python -m pytest tests/unit/test_sql_execution.py::TestCreateRetrySession::test_retry_session_has_correct_config -v 2>&1 | head -50Repository: deepnote/deepnote-toolkit
Length of output: 1964
🏁 Script executed:
#!/bin/bash
# Check if the test file exists and read the relevant section
cd "$(git rev-parse --show-toplevel)"
if [ -f tests/unit/test_sql_execution.py ]; then
echo "=== Test file content around lines 1026-1040 ==="
sed -n '1026,1040p' tests/unit/test_sql_execution.py
else
echo "Test file not found"
fiRepository: deepnote/deepnote-toolkit
Length of output: 882
🏁 Script executed:
#!/bin/bash
# Find the _create_retry_session function
cd "$(git rev-parse --show-toplevel)"
fd -e py -x grep -l "_create_retry_session" {} \;Repository: deepnote/deepnote-toolkit
Length of output: 143
🏁 Script executed:
#!/bin/bash
# Read the _create_retry_session implementation
cd "$(git rev-parse --show-toplevel)"
rg -A 30 "def _create_retry_session" deepnote_toolkit/sql/sql_execution.pyRepository: deepnote/deepnote-toolkit
Length of output: 1056
Add allowed_methods parameter to the Retry configuration.
Line 1040's assertion expects allowed_methods to contain "POST", but the current _create_retry_session() implementation doesn't set this parameter. The Retry object should include allowed_methods=["DELETE", "GET", "HEAD", "OPTIONS", "POST", "PUT", "TRACE"] to enable retries on POST requests as the docstring claims.
🧰 Tools
🪛 Ruff (0.15.7)
[warning] 1037-1037: Use a regular assert instead of unittest-style assertEqual
Replace assertEqual(...) with assert ...
(PT009)
[warning] 1038-1038: Use a regular assert instead of unittest-style assertEqual
Replace assertEqual(...) with assert ...
(PT009)
[warning] 1039-1039: Use a regular assert instead of unittest-style assertEqual
Replace assertEqual(...) with assert ...
(PT009)
[warning] 1040-1040: Use a regular assert instead of unittest-style assertIn
Replace assertIn(...) with assert ...
(PT009)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit/test_sql_execution.py` around lines 1026 - 1040, The Retry
configuration in _create_retry_session is missing the allowed_methods parameter
so POST retries aren't enabled; update the Retry(...) construction inside the
_create_retry_session function to include
allowed_methods=["DELETE","GET","HEAD","OPTIONS","POST","PUT","TRACE"] (or an
equivalent set that includes "POST"), then ensure that same Retry instance is
used when mounting the HTTP/HTTPS adapters so
adapter.max_retries.allowed_methods contains "POST".
2 participants
_create_retry_sessionto create a requests session with retry capabilities for handling 5xx errors on POST requests._generate_temporary_credentialsand_get_federated_auth_credentialsto use the new retry session for making requests.Summary by CodeRabbit
New Features
Tests