Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions deepnote_toolkit/sql/sql_execution.py
Original file line number Diff line number Diff line change
Expand Up @@ -626,6 +626,22 @@ def _cancel_cursor(cursor: "DBAPICursor") -> None:
pass # Best effort, ignore all errors


def _extract_rls_user_email(query):
"""Extract the Deepnote user email from the audit comment embedded in *query*.

Deepnote appends an audit comment such as ``/* {"user_email":"a@b.com",...} */`` to
every SQL block. We reuse that identity for row-level security. Returns "" when no
email is present (e.g. anonymous app viewers) so that the policy denies all rows.

The value is restricted to a safe character set so it can be embedded directly into a
``SET`` statement without risk of SQL injection, regardless of the caller.
"""
match = re.search(r'"user_email"\s*:\s*"([^"]+)"', query or "")
if not match:
return ""
return re.sub(r"[^A-Za-z0-9.+@_-]", "", match.group(1))

Comment on lines +629 to +643

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Add type hints.

Function lacks explicit type annotations for parameter and return value.

-def _extract_rls_user_email(query):
+def _extract_rls_user_email(query: 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.

Suggested change
def _extract_rls_user_email(query):
"""Extract the Deepnote user email from the audit comment embedded in *query*.
Deepnote appends an audit comment such as ``/* {"user_email":"a@b.com",...} */`` to
every SQL block. We reuse that identity for row-level security. Returns "" when no
email is present (e.g. anonymous app viewers) so that the policy denies all rows.
The value is restricted to a safe character set so it can be embedded directly into a
``SET`` statement without risk of SQL injection, regardless of the caller.
"""
match = re.search(r'"user_email"\s*:\s*"([^"]+)"', query or "")
if not match:
return ""
return re.sub(r"[^A-Za-z0-9.+@_-]", "", match.group(1))
def _extract_rls_user_email(query: str) -> str:
"""Extract the Deepnote user email from the audit comment embedded in *query*.
Deepnote appends an audit comment such as ``/* {"user_email":"a@b.com",...} */`` to
every SQL block. We reuse that identity for row-level security. Returns "" when no
email is present (e.g. anonymous app viewers) so that the policy denies all rows.
The value is restricted to a safe character set so it can be embedded directly into a
``SET`` statement without risk of SQL injection, regardless of the caller.
"""
match = re.search(r'"user_email"\s*:\s*"([^"]+)"', query or "")
if not match:
return ""
return re.sub(r"[^A-Za-z0-9.+@_-]", "", match.group(1))
🧰 Tools
🪛 Ruff (0.15.17)

[warning] 629-629: Missing return type annotation for private function _extract_rls_user_email

(ANN202)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@deepnote_toolkit/sql/sql_execution.py` around lines 629 - 643, The function
_extract_rls_user_email lacks explicit type annotations for its parameter and
return value. Add a type hint for the query parameter to indicate it accepts an
optional string (using Optional[str] or str | None based on the `query or ""`
usage pattern), and add a return type hint of str since the function always
returns a string.

Sources: Coding guidelines, Linters/SAST tools


def _execute_sql_on_engine(engine, query, bind_params):
"""Run *query* on *engine* and return a DataFrame.

Expand All @@ -639,6 +655,7 @@ def _execute_sql_on_engine(engine, query, bind_params):

import pandas as pd
from sqlalchemy import __version__ as sqlalchemy_version
from sqlalchemy import text

from deepnote_toolkit.config import get_config

Expand All @@ -656,6 +673,16 @@ def _execute_sql_on_engine(engine, query, bind_params):
)

with engine.begin() as connection:
# Row-level security: expose the Deepnote user identity (already carried in the
# audit comment) as a Snowflake session variable so row access policies can read
# it via GETVARIABLE('DEEPNOTE_USER_EMAIL'). Always set it on every Snowflake
# query — session variables persist on pooled connections, so an unset/stale
# value could leak another user's rows. Empty value => policy matches nothing
# (deny by default for anonymous viewers).
if engine.dialect.name == "snowflake":
rls_user_email = _extract_rls_user_email(query)
connection.execute(text(f"SET DEEPNOTE_USER_EMAIL = '{rls_user_email}'"))

# For pandas 2.2+ with SQLAlchemy < 2.0, use raw DBAPI connection
if needs_raw_connection:
tracking_connection = CursorTrackingDBAPIConnection(connection.connection)
Expand Down
Loading