fix(big-query): Fix BigQuery federated auth token refresh

[tkislan]

Summary by CodeRabbit

• Improvements

  • Enhanced authentication and IAM parameter handling in SQL execution for improved reliability and consistency.

• Tests

  • Added comprehensive unit tests validating federated authentication credentials verification and response handling.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Authentication parameter handling is reorganized in SQL execution. Federated authentication and IAM parameter processing now occur at the start of execute_sql_with_connection_json() instead of being called within _query_data_source(). This ensures credential augmentation completes before URL construction and SSH handling. A new test validates federated authentication credential retrieval, including URL composition, request formatting, and response parsing.

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	Title directly matches the main change: moving IAM and federated-auth handling earlier in the execution flow for BigQuery federated auth token refresh.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

---

📜 Recent 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.

📥 Commits

Reviewing files that changed from the base of the PR and between 10ab3d7 and 9316523.

📒 Files selected for processing (2)

• deepnote_toolkit/sql/sql_execution.py
• tests/unit/test_sql_execution.py

🧰 Additional context used

📓 Path-based instructions (3)

**/*.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 use Optional[T] for parameters that can be None (not T = 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:

• tests/unit/test_sql_execution.py
• deepnote_toolkit/sql/sql_execution.py

**/test_*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Name test files with test_*.py pattern

Files:

• tests/unit/test_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

🧬 Code graph analysis (1)

tests/unit/test_sql_execution.py (1)

deepnote_toolkit/sql/sql_execution.py (1)

• _get_federated_auth_credentials (271-290)

🪛 Ruff (0.14.11)

tests/unit/test_sql_execution.py

636-636: Use a regular assert instead of unittest-style assertEqual

Replace assertEqual(...) with assert ...

(PT009)

---

637-637: Use a regular assert instead of unittest-style assertEqual

Replace assertEqual(...) with assert ...

(PT009)

⏰ 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). (8)

• GitHub Check: Test - Python 3.11
• GitHub Check: Test - Python 3.10
• GitHub Check: Test - Python 3.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.13
• GitHub Check: Build and push artifacts for Python 3.11
• GitHub Check: Build and push artifacts for Python 3.12

🔇 Additional comments (2)

deepnote_toolkit/sql/sql_execution.py (1)

118-120: LGTM - handlers moved to correct position.

Credential augmentation now completes before URL construction and SSH handling. Clean refactor.

tests/unit/test_sql_execution.py (1)

592-637: Good coverage of federated auth flow.

Test validates URL construction, headers composition, status check, and response parsing. Solid.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

📦 Python package built successfully!

• Version: 1.2.0.dev3+cd24a36
• Wheel: deepnote_toolkit-1.2.0.dev3+cd24a36-py3-none-any.whl
• Install:

  pip install "deepnote-toolkit @ https://deepnote-staging-runtime-artifactory.s3.amazonaws.com/deepnote-toolkit-packages/1.2.0.dev3%2Bcd24a36/deepnote_toolkit-1.2.0.dev3%2Bcd24a36-py3-none-any.whl"

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.27%. Comparing base (10ab3d7) to head (9316523).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #54      +/-   ##
==========================================
+ Coverage   73.14%   73.27%   +0.13%     
==========================================
  Files          93       93              
  Lines        5194     5194              
  Branches      758      758              
==========================================
+ Hits         3799     3806       +7     
+ Misses       1151     1144       -7     
  Partials      244      244              

Flag	Coverage Δ
combined	73.27% <100.00%> (+0.13%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

deepnote_toolkit/sql/sql_execution.py (2)

337-343: TypeError not caught for BigQuery.

If sql_alchemy_dict["params"] is None, line 339 raises TypeError, not KeyError. The except block won't catch it.

Proposed fix

     elif federated_auth.integrationType == "big-query":
         try:
+            if sql_alchemy_dict.get("params") is None:
+                raise KeyError("params")
             sql_alchemy_dict["params"]["access_token"] = federated_auth.accessToken
         except KeyError:
             logger.exception(
                 "Invalid federated auth params, try updating toolkit version"
             )

---

280-282: Use dictionary unpacking for headers instead of in-place mutation.

get_project_auth_headers() returns a fresh dict on each call, so the mutation is safe. However, follow the dictionary unpacking pattern from the coding guidelines:

headers = {
    "UserPodAuthContextToken": user_pod_auth_context_token,
    **get_project_auth_headers(),
}

tests/unit/test_sql_execution.py (1)

905-948: Test doesn't cover TypeError case.

Test test_federated_auth_params_bigquery_missing_params validates when params key is absent. But if params is None, production code raises TypeError, not KeyError. Consider adding a test case for "params": None.

📜 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.

📥 Commits

Reviewing files that changed from the base of the PR and between 10ab3d7 and 9316523.

📒 Files selected for processing (2)

• deepnote_toolkit/sql/sql_execution.py
• tests/unit/test_sql_execution.py

🧰 Additional context used

📓 Path-based instructions (3)

**/*.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 use Optional[T] for parameters that can be None (not T = 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:

• tests/unit/test_sql_execution.py
• deepnote_toolkit/sql/sql_execution.py

**/test_*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Name test files with test_*.py pattern

Files:

• tests/unit/test_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

🧬 Code graph analysis (1)

tests/unit/test_sql_execution.py (1)

deepnote_toolkit/sql/sql_execution.py (1)

• _get_federated_auth_credentials (271-290)

🪛 Ruff (0.14.11)

tests/unit/test_sql_execution.py

636-636: Use a regular assert instead of unittest-style assertEqual

Replace assertEqual(...) with assert ...

(PT009)

---

637-637: Use a regular assert instead of unittest-style assertEqual

Replace assertEqual(...) with assert ...

(PT009)

⏰ 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). (8)

• GitHub Check: Test - Python 3.11
• GitHub Check: Test - Python 3.10
• GitHub Check: Test - Python 3.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.13
• GitHub Check: Build and push artifacts for Python 3.11
• GitHub Check: Build and push artifacts for Python 3.12

🔇 Additional comments (2)

deepnote_toolkit/sql/sql_execution.py (1)

118-120: LGTM - handlers moved to correct position.

Credential augmentation now completes before URL construction and SSH handling. Clean refactor.

tests/unit/test_sql_execution.py (1)

592-637: Good coverage of federated auth flow.

Test validates URL construction, headers composition, status check, and response parsing. Solid.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

[deepnote-bot]

🚀 Review App Deployment Started

📝 Description	🌐 Link / Info
🌍 Review application	ra-54
🔑 Sign-in URL	Click to sign-in
📊 Application logs	View logs
🔄 Actions	Click to redeploy
🚀 ArgoCD deployment	View deployment
⏰ Last deployed	2026-01-14 14:06:23 (UTC)
📜 Deployed commit	e477fd817a252a2f8cdeced74602f8e2ef2317a6
🛠️ Toolkit version	cd24a36