feat: Add support to refresh federated auth access token #46
Conversation
📝 WalkthroughPre-merge checks❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
Comment |
|
📦 Python package built successfully!
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #46 +/- ##
=======================================
Coverage 73.04% 73.04%
=======================================
Files 93 93
Lines 5149 5187 +38
Branches 754 758 +4
=======================================
+ Hits 3761 3789 +28
- Misses 1144 1153 +9
- Partials 244 245 +1
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.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
deepnote_toolkit/sql/sql_execution.py (2)
317-319: Uselogger.exceptionfor automatic traceback capture.Per static analysis hint TRY400 and coding guidelines.
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
326-329: Missing KeyError protection for nested dict access.Lines 327-329 assume
sql_alchemy_dict["params"]["connect_args"]["http_headers"]exists. Will raiseKeyErrorif the dict structure is incomplete.🔎 Proposed fix
if federated_auth.integrationType == "trino": - sql_alchemy_dict["params"]["connect_args"]["http_headers"][ - "Authorization" - ] = f"Bearer {federated_auth.accessToken}" + params = sql_alchemy_dict.setdefault("params", {}) + connect_args = params.setdefault("connect_args", {}) + http_headers = connect_args.setdefault("http_headers", {}) + http_headers["Authorization"] = f"Bearer {federated_auth.accessToken}" elif federated_auth.integrationType == "big-query": - sql_alchemy_dict["params"]["access_token"] = federated_auth.accessToken + params = sql_alchemy_dict.setdefault("params", {}) + params["access_token"] = federated_auth.accessToken
📜 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
317-319: 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: Build and push artifacts for Python 3.9
- GitHub Check: Build and push artifacts for Python 3.10
- GitHub Check: Build and push artifacts for Python 3.12
- GitHub Check: Build and push artifacts for Python 3.13
- GitHub Check: Build and push artifacts for Python 3.11
- GitHub Check: Typecheck - 3.13
- GitHub Check: Test - Python 3.11
- GitHub Check: Test - Python 3.9
- GitHub Check: Test - Python 3.10
🔇 Additional comments (4)
deepnote_toolkit/sql/sql_execution.py (4)
4-4: LGTM: Imports and logger setup.Standard imports and logger initialization for new federated auth functionality.
Also applies to: 8-8, 19-19, 39-40
260-260: LGTM: HTTP error handling added.Good addition of
raise_for_status()for early failure detection.
289-304: LGTM: Clean IAM credential handling.Proper early return, clear logic, follows guidelines.
442-442: LGTM: Clean integration of credential handlers.Proper placement and ordering of new helper functions.
Also applies to: 444-444
| class IntegrationFederatedAuthParams(BaseModel): | ||
| integrationId: str | ||
| authContextToken: str | ||
|
|
||
|
|
||
| class FederatedAuthResponseData(BaseModel): | ||
| integrationType: str | ||
| accessToken: str | ||
|
|
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider snake_case attributes with camelCase aliases.
The model fields use camelCase, which matches the API but violates Python conventions. For better ergonomics, use snake_case attributes with Pydantic aliases:
from pydantic import BaseModel, Field
class IntegrationFederatedAuthParams(BaseModel):
integration_id: str = Field(alias="integrationId")
auth_context_token: str = Field(alias="authContextToken")
class FederatedAuthResponseData(BaseModel):
integration_type: str = Field(alias="integrationType")
access_token: str = Field(alias="accessToken")This allows idiomatic Python access (federated_auth.access_token) while preserving API compatibility.
🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 42 to 50, the Pydantic
models use camelCase field names which break Python naming conventions; change
fields to snake_case attributes and add Pydantic aliases that map to the
existing camelCase keys (use Field(alias="...") for each field) so external API
compatibility remains while callers can use idiomatic snake_case attribute
access; ensure models still accept population by field name and by alias if
needed (set Config allow_population_by_field_name=True or equivalent) and update
any downstream code that references the old attribute names.
| def _get_federated_auth_credentials( | ||
| integration_id: str, user_pod_auth_context_token: str | ||
| ) -> FederatedAuthResponseData: | ||
| """Get federated auth credentials for the given integration ID and user pod auth context token.""" | ||
|
|
||
| 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 | ||
|
|
There was a problem hiding this comment.
Add error handling for JSON decode and validation.
response.json() and model_validate() can raise exceptions that should be caught and logged with context.
🔎 Proposed fix
def _get_federated_auth_credentials(
integration_id: str, user_pod_auth_context_token: str
) -> FederatedAuthResponseData:
"""Get federated auth credentials for the given integration ID and user pod auth context token."""
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()
+ try:
- data = FederatedAuthResponseData.model_validate(response.json())
+ data = FederatedAuthResponseData.model_validate(response.json())
+ except (ValueError, ValidationError) as e:
+ logger.error(
+ "Failed to parse federated auth response from %s: %s",
+ url,
+ response.text,
+ exc_info=e,
+ )
+ raise
return dataCommittable suggestion skipped: line range outside the PR's diff.
Summary by CodeRabbit
New Features
Refactor
Chores
Tests
✏️ Tip: You can customize this high-level summary in your review settings.