feat: Add support to refresh federated auth access token #46
Conversation
📝 WalkthroughWalkthroughAdds federated authentication support for SQL execution by introducing two Pydantic models (IntegrationFederatedAuthParams, FederatedAuthResponseData) and three helpers. _get_federated_auth_credentials(...) requests a federated token from an API and returns structured data. _handle_iam_params(...) and _handle_federated_auth_params(...) update sql_alchemy_dict in-place to apply IAM or federated credentials. _query_data_source(...) now calls these helpers, includes response.raise_for_status() after HTTP calls, and logs/handles validation errors when federated-auth params are invalid. Sequence Diagram(s)sequenceDiagram
participant Query as _query_data_source
participant IAM as _handle_iam_params
participant FED as _handle_federated_auth_params
participant TokenAPI as Federated Auth API
participant DB as Database / SQLAlchemy
Query->>IAM: inspect sql_alchemy_dict for IAM params
IAM-->>Query: apply IAM creds into sql_alchemy_dict
Query->>FED: inspect sql_alchemy_dict for federated params
FED->>TokenAPI: GET token (integration_id, user_pod_auth_context_token)
TokenAPI-->>FED: 200 + token / error (HTTP status surfaced)
FED-->>Query: apply federated token into sql_alchemy_dict or log/skip on validation error
Query->>DB: create connection using sql_alchemy_dict and execute query
DB-->>Query: return results or error
Pre-merge checks❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Disabled knowledge base sources:
📒 Files selected for processing (1)
🧰 Additional context used📓 Path-based instructions (2)**/*.py📄 CodeRabbit inference engine (.cursorrules)
Files:
deepnote_toolkit/**/*.py📄 CodeRabbit inference engine (.cursorrules)
Files:
🪛 Ruff (0.14.8)deepnote_toolkit/sql/sql_execution.py315-317: Use Replace with (TRY400) ⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
🔇 Additional comments (4)
Comment |
|
📦 Python package built successfully!
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #46 +/- ##
==========================================
- Coverage 73.04% 72.80% -0.24%
==========================================
Files 93 93
Lines 5149 5185 +36
Branches 754 757 +3
==========================================
+ Hits 3761 3775 +14
- Misses 1144 1164 +20
- Partials 244 246 +2
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
📜 Review details
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Disabled knowledge base sources:
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
deepnote_toolkit/sql/sql_execution.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (.cursorrules)
**/*.py: Write clean, readable Python code following PEP 8 style guide
Use type hints with Optional[T] for parameters that can be None (not T = None)
Use explicit type hints for function parameters and return values
Maximum line length: 88 characters (Black default)
Use f-strings instead of .format() for string formatting
Use pathlib.Path for file path operations instead of os.path
Use black for code formatting
Use isort for import sorting (black profile)
Use flake8 for linting
Use early returns to reduce nesting and extract common checks into variables for readability
Use snake_case for variable and function names
Use PascalCase for class names
Use snake_case for file names
Support Python versions 3.9, 3.10, 3.11, 3.12, and 3.13
**/*.py: Follow PEP 8 with Black formatting (line length: 88)
Use isort with Black profile for import sorting
Use type hints consistently
Use docstrings for all functions/classes
Use f-strings instead of .format() for string formatting
Use pathlib.Path for file path operations instead of os.path
Always useOptional[T]for parameters that can be None (notT = None)
Use explicit type hints for function parameters and return values
Use snake_case for files, functions, and variables
Use PascalCase for classes
Use appropriate exception types with context logging for error handling
Handle Jupyter/IPython specific exceptions properly
Use early returns to reduce nesting and extract common checks into variables for readability
Use dictionary unpacking for headers (e.g.,headers = {"Content-Type": "application/json", **auth_headers})
Use space-separated format for CLI arguments (e.g.,--port 8080)
Files:
deepnote_toolkit/sql/sql_execution.py
deepnote_toolkit/**/*.py
📄 CodeRabbit inference engine (.cursorrules)
deepnote_toolkit/**/*.py: Use dictionary unpacking for headers: headers = {"Content-Type": "application/json", **auth_headers}
Use appropriate exception types, log errors with context, and handle Jupyter/IPython specific exceptions properly
Document functions and classes with docstrings
Files:
deepnote_toolkit/sql/sql_execution.py
🪛 Ruff (0.14.8)
deepnote_toolkit/sql/sql_execution.py
310-312: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
311-311: Logging statement uses f-string
(G004)
319-319: Cannot use match statement on Python 3.9 (syntax was added in Python 3.10)
(invalid-syntax)
328-330: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
- GitHub Check: Test - Python 3.12
- GitHub Check: Typecheck - 3.13
- GitHub Check: Test - Python 3.10
- GitHub Check: Test - Python 3.11
- GitHub Check: Build and push artifacts for Python 3.9
- GitHub Check: Typecheck - 3.9
- GitHub Check: Build and push artifacts for Python 3.10
- GitHub Check: Build and push artifacts for Python 3.11
- GitHub Check: Build and push artifacts for Python 3.12
- GitHub Check: Build and push artifacts for Python 3.13
🔇 Additional comments (4)
deepnote_toolkit/sql/sql_execution.py (4)
10-13:Neverimport looks correct for version compatibility.Proper conditional import for Python 3.9 support.
48-56: Pydantic models defined correctly.Clean structure for federated auth validation.
321-323: Potential KeyError if nested keys don't exist.If
params,connect_args, orhttp_headersare missing, this will raise KeyError.Verify that callers always provide the expected structure, or add defensive checks:
sql_alchemy_dict.setdefault("params", {}).setdefault("connect_args", {}).setdefault("http_headers", {})["Authorization"] = f"Bearer {access_token}"
432-434: Clean refactor to helper functions.Inline IAM/federated logic is now properly encapsulated.
| def _get_federated_auth_credentials(integration_id: str, user_id: str) -> str: | ||
| url = get_absolute_userpod_api_url( | ||
| f"integrations/federated-auth-token/{integration_id}" | ||
| ) | ||
|
|
||
| # Add project credentials in detached mode | ||
| headers = get_project_auth_headers() | ||
|
|
||
| response = requests.post(url, json={"userId": user_id}, timeout=10, headers=headers) | ||
|
|
||
| data = FederatedAuthResponseData.model_validate_json(response.json()) | ||
|
|
||
| return data.accessToken | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Missing docstring and error handling for HTTP failures.
Per coding guidelines, functions need docstrings. Also, response.json() will fail silently on non-2xx responses.
def _get_federated_auth_credentials(integration_id: str, user_id: str) -> str:
+ """Fetch federated auth access token from the userpod API."""
url = get_absolute_userpod_api_url(
f"integrations/federated-auth-token/{integration_id}"
)
# Add project credentials in detached mode
headers = get_project_auth_headers()
response = requests.post(url, json={"userId": user_id}, timeout=10, headers=headers)
+ response.raise_for_status()
- data = FederatedAuthResponseData.model_validate_json(response.json())
+ data = FederatedAuthResponseData.model_validate(response.json())
return data.accessToken🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 271 to 284, add a concise
docstring describing the function purpose, parameters (integration_id, user_id)
and return value (access token string), and implement robust HTTP/error
handling: call response.raise_for_status() after the POST to fail fast on
non-2xx responses, wrap response.json() and
FederatedAuthResponseData.model_validate_json(...) in a try/except to catch
JSONDecodeError and validation exceptions and raise or log a clear exception
with context (including status code and response.text), and ensure the function
still returns the access token string on success.
| except ValidationError as e: | ||
| logger.error( | ||
| f"Invalid federated auth params, try updating toolkit version: {e}" | ||
| ) | ||
| return |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Use logger.exception to include traceback.
Per static analysis hint TRY400. Also, avoid f-strings in logging calls.
except ValidationError as e:
- logger.error(
- f"Invalid federated auth params, try updating toolkit version: {e}"
- )
+ logger.exception("Invalid federated auth params, try updating toolkit version")
return🧰 Tools
🪛 Ruff (0.14.8)
310-312: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
311-311: Logging statement uses f-string
(G004)
🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 309 to 313, the except
block currently logs the ValidationError using logger.error with an f-string
which omits the traceback and misuses logging formatting; replace that call with
logger.exception and use structured logging parameters (e.g.,
logger.exception("Invalid federated auth params, try updating toolkit version:
%s", e)) so the traceback is included and avoid f-string interpolation in the
logging call.
|
🚀 Review App Deployment Started
|
…ype hints - Refactored `_get_federated_auth_credentials` to return `FederatedAuthResponseData` instead of just the access token. - Updated `IntegrationFederatedAuthParams` to include `userPodAuthContextToken` and changed the handling of integration types. - Improved error logging for unsupported integration types. - Cleaned up imports and ensured consistent use of type hints.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
deepnote_toolkit/sql/sql_execution.py (2)
265-278: Still missing docstring and error handling.Per previous review, this function needs:
- A docstring describing purpose, parameters, and return value
response.raise_for_status()to fail fast on non-2xx responses- Error handling for
response.json()and validation failuresWithout these, silent failures will cause confusing downstream errors.
🔎 Proposed fix
def _get_federated_auth_credentials(integration_id: str, user_pod_auth_context_token: str) -> FederatedAuthResponseData: + """Fetch federated auth credentials from the userpod API. + + Args: + integration_id: The integration ID to fetch credentials for + user_pod_auth_context_token: Auth context token for the user pod + + Returns: + FederatedAuthResponseData containing access token and integration type + """ url = get_absolute_userpod_api_url( f"integrations/federated-auth-token/{integration_id}" ) # Add project credentials in detached mode headers = get_project_auth_headers() headers["UserPodAuthContextToken"] = user_pod_auth_context_token response = requests.post(url, timeout=10, headers=headers) + response.raise_for_status() data = FederatedAuthResponseData.model_validate(response.json()) return data
309-311: Uselogger.exceptionfor automatic traceback.Per static analysis and previous review, replace
logger.error(..., exc_info=e)withlogger.exception(...)for cleaner traceback inclusion.🔎 Proposed fix
except ValidationError as e: - logger.error( - "Invalid federated auth params, try updating toolkit version:", exc_info=e - ) + logger.exception("Invalid federated auth params, try updating toolkit version") return
📜 Review details
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Disabled knowledge base sources:
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
deepnote_toolkit/sql/sql_execution.py(5 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (.cursorrules)
**/*.py: Write clean, readable Python code following PEP 8 style guide
Use type hints with Optional[T] for parameters that can be None (not T = None)
Use explicit type hints for function parameters and return values
Maximum line length: 88 characters (Black default)
Use f-strings instead of .format() for string formatting
Use pathlib.Path for file path operations instead of os.path
Use black for code formatting
Use isort for import sorting (black profile)
Use flake8 for linting
Use early returns to reduce nesting and extract common checks into variables for readability
Use snake_case for variable and function names
Use PascalCase for class names
Use snake_case for file names
Support Python versions 3.9, 3.10, 3.11, 3.12, and 3.13
**/*.py: Follow PEP 8 with Black formatting (line length: 88)
Use isort with Black profile for import sorting
Use type hints consistently
Use docstrings for all functions/classes
Use f-strings instead of .format() for string formatting
Use pathlib.Path for file path operations instead of os.path
Always useOptional[T]for parameters that can be None (notT = None)
Use explicit type hints for function parameters and return values
Use snake_case for files, functions, and variables
Use PascalCase for classes
Use appropriate exception types with context logging for error handling
Handle Jupyter/IPython specific exceptions properly
Use early returns to reduce nesting and extract common checks into variables for readability
Use dictionary unpacking for headers (e.g.,headers = {"Content-Type": "application/json", **auth_headers})
Use space-separated format for CLI arguments (e.g.,--port 8080)
Files:
deepnote_toolkit/sql/sql_execution.py
deepnote_toolkit/**/*.py
📄 CodeRabbit inference engine (.cursorrules)
deepnote_toolkit/**/*.py: Use dictionary unpacking for headers: headers = {"Content-Type": "application/json", **auth_headers}
Use appropriate exception types, log errors with context, and handle Jupyter/IPython specific exceptions properly
Document functions and classes with docstrings
Files:
deepnote_toolkit/sql/sql_execution.py
🪛 Ruff (0.14.8)
deepnote_toolkit/sql/sql_execution.py
309-311: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)
- GitHub Check: Test - Python 3.11
- GitHub Check: Test - Python 3.9
- GitHub Check: Build and push artifacts for Python 3.10
- GitHub Check: Test - Python 3.10
- GitHub Check: Build and push artifacts for Python 3.13
- GitHub Check: Build and push artifacts for Python 3.12
- GitHub Check: Build and push artifacts for Python 3.11
- GitHub Check: Build and push artifacts for Python 3.9
- GitHub Check: Typecheck - 3.13
🔇 Additional comments (3)
deepnote_toolkit/sql/sql_execution.py (3)
4-4: LGTM.Imports and model setup are clean. The camelCase field names in Pydantic models align with typical API contracts.
Also applies to: 8-8, 19-19, 39-49
281-295: LGTM.Docstring added and variable naming fixed per previous review.
429-431: LGTM.Clean refactoring to use the new helper functions.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
deepnote_toolkit/sql/sql_execution.py (1)
309-311: Fix logging: useexc_info=Truenotexc_info=e.Should be
logger.exception(...)orlogger.error(..., exc_info=True).Based on static analysis hint TRY400.
🔎 Proposed fix
except ValidationError as e: - logger.error( - "Invalid federated auth params, try updating toolkit version:", exc_info=e - ) + logger.exception("Invalid federated auth params, try updating toolkit version") return
📜 Review details
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Disabled knowledge base sources:
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
deepnote_toolkit/sql/sql_execution.py(5 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (.cursorrules)
**/*.py: Write clean, readable Python code following PEP 8 style guide
Use type hints with Optional[T] for parameters that can be None (not T = None)
Use explicit type hints for function parameters and return values
Maximum line length: 88 characters (Black default)
Use f-strings instead of .format() for string formatting
Use pathlib.Path for file path operations instead of os.path
Use black for code formatting
Use isort for import sorting (black profile)
Use flake8 for linting
Use early returns to reduce nesting and extract common checks into variables for readability
Use snake_case for variable and function names
Use PascalCase for class names
Use snake_case for file names
Support Python versions 3.9, 3.10, 3.11, 3.12, and 3.13
**/*.py: Follow PEP 8 with Black formatting (line length: 88)
Use isort with Black profile for import sorting
Use type hints consistently
Use docstrings for all functions/classes
Use f-strings instead of .format() for string formatting
Use pathlib.Path for file path operations instead of os.path
Always useOptional[T]for parameters that can be None (notT = None)
Use explicit type hints for function parameters and return values
Use snake_case for files, functions, and variables
Use PascalCase for classes
Use appropriate exception types with context logging for error handling
Handle Jupyter/IPython specific exceptions properly
Use early returns to reduce nesting and extract common checks into variables for readability
Use dictionary unpacking for headers (e.g.,headers = {"Content-Type": "application/json", **auth_headers})
Use space-separated format for CLI arguments (e.g.,--port 8080)
Files:
deepnote_toolkit/sql/sql_execution.py
deepnote_toolkit/**/*.py
📄 CodeRabbit inference engine (.cursorrules)
deepnote_toolkit/**/*.py: Use dictionary unpacking for headers: headers = {"Content-Type": "application/json", **auth_headers}
Use appropriate exception types, log errors with context, and handle Jupyter/IPython specific exceptions properly
Document functions and classes with docstrings
Files:
deepnote_toolkit/sql/sql_execution.py
🪛 Ruff (0.14.8)
deepnote_toolkit/sql/sql_execution.py
309-311: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)
- GitHub Check: Test - Python 3.11
- GitHub Check: Test - Python 3.10
- GitHub Check: Typecheck - 3.13
- GitHub Check: Test - Python 3.9
- GitHub Check: Build and push artifacts for Python 3.12
- GitHub Check: Build and push artifacts for Python 3.9
- GitHub Check: Build and push artifacts for Python 3.13
- GitHub Check: Build and push artifacts for Python 3.10
- GitHub Check: Build and push artifacts for Python 3.11
🔇 Additional comments (1)
deepnote_toolkit/sql/sql_execution.py (1)
429-431: Clean integration of credential handlers.The in-place mutation approach keeps the flow straightforward.
| def _get_federated_auth_credentials(integration_id: str, user_pod_auth_context_token: str) -> FederatedAuthResponseData: | ||
| url = get_absolute_userpod_api_url( | ||
| f"integrations/federated-auth-token/{integration_id}" | ||
| ) | ||
|
|
||
| # Add project credentials in detached mode | ||
| headers = get_project_auth_headers() | ||
| headers["UserPodAuthContextToken"] = user_pod_auth_context_token | ||
|
|
||
| response = requests.post(url, timeout=10, headers=headers) | ||
|
|
||
| data = FederatedAuthResponseData.model_validate(response.json()) | ||
|
|
||
| return data |
There was a problem hiding this comment.
Missing HTTP error handling and docstring.
- No
response.raise_for_status()call—HTTP errors will cause silent failures or JSON decode errors. - Docstring required per coding guidelines.
🔎 Proposed fix
def _get_federated_auth_credentials(integration_id: str, user_pod_auth_context_token: str) -> FederatedAuthResponseData:
+ """Fetch federated auth credentials from userpod API.
+
+ Args:
+ integration_id: Integration identifier
+ user_pod_auth_context_token: Auth context token
+
+ Returns:
+ Federated auth response containing access token and integration type
+ """
url = get_absolute_userpod_api_url(
f"integrations/federated-auth-token/{integration_id}"
)
# Add project credentials in detached mode
headers = get_project_auth_headers()
headers["UserPodAuthContextToken"] = user_pod_auth_context_token
response = requests.post(url, timeout=10, headers=headers)
+ response.raise_for_status()
data = FederatedAuthResponseData.model_validate(response.json())
return data🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 265 to 278, the helper
_get_federated_auth_credentials lacks a docstring and does not handle HTTP
errors; add a short docstring describing parameters and return type, call
response.raise_for_status() immediately after the POST to raise on non-2xx
responses, and wrap response.json()/model_validate in a try/except to raise a
clear exception (including response.status_code and response.text) on JSON
decode or validation errors so failures are not silent.
| if federated_auth.integrationType == "trino": | ||
| sql_alchemy_dict["params"]["connect_args"]["http_headers"][ | ||
| "Authorization" | ||
| ] = f"Bearer {federated_auth.accessToken}" | ||
| elif federated_auth.integrationType == "big-query": | ||
| sql_alchemy_dict["params"]["access_token"] = federated_auth.accessToken |
There was a problem hiding this comment.
Missing KeyError handling for nested dict access.
Lines 319-321 assume sql_alchemy_dict["params"]["connect_args"]["http_headers"] exists. Will raise KeyError if structure is incomplete.
🔎 Proposed fix
if federated_auth.integrationType == "trino":
- sql_alchemy_dict["params"]["connect_args"]["http_headers"][
- "Authorization"
- ] = f"Bearer {federated_auth.accessToken}"
+ try:
+ sql_alchemy_dict["params"]["connect_args"]["http_headers"][
+ "Authorization"
+ ] = f"Bearer {federated_auth.accessToken}"
+ except KeyError:
+ logger.error(
+ "Missing required connection structure for Trino federated auth"
+ )
+ return
elif federated_auth.integrationType == "big-query":
sql_alchemy_dict["params"]["access_token"] = federated_auth.accessToken🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 318 to 323 the code
assumes nested keys sql_alchemy_dict["params"]["connect_args"]["http_headers"]
exist and will raise KeyError if the dict structure is incomplete; update the
code to defensively ensure each nested mapping exists (e.g., use dict.setdefault
or check and assign empty dicts for "params", "connect_args" and "http_headers")
before assigning the Authorization header, and likewise ensure "params" exists
before setting "access_token" for big-query so the assignment never raises
KeyError.
Summary by CodeRabbit
New Features
Refactor
Chores
✏️ Tip: You can customize this high-level summary in your review settings.