Skip to content
Merged
Show file tree
Hide file tree
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
21 changes: 20 additions & 1 deletion superset/mcp_service/sql_lab/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,26 @@ class OpenSqlLabRequest(BaseModel):
description="SQL to pre-populate in the editor",
validation_alias=AliasChoices("sql", "query"),
)
title: str | None = Field(None, description="Title for the SQL Lab tab/query")
title: str | None = Field(
None,
description=(
"Title for the SQL Lab tab. Generate a succinct, descriptive label "
"(roughly 3-6 words) from the conversation context or the query "
"intent — e.g. 'Top customers by revenue Q3' — instead of leaving "
"the tab untitled. Avoid generic names like 'Untitled Query'."
),
max_length=256,
)

@field_validator("title")
@classmethod
def title_strip_or_none(cls, v: str | None) -> str | None:
# Whitespace-only would render as a blank tab label; fall back to
# SQL Lab's default "Untitled Query N" naming instead.
if v is None:
return None
stripped = v.strip()
return stripped or None


class SqlLabResponse(_SchemaFieldNormalizer):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@

logger = logging.getLogger(__name__)

SQL_LAB_QUERY_PARAMS_TO_SANITIZE = frozenset({"sql", "title"})
SQL_LAB_QUERY_PARAMS_TO_SANITIZE = frozenset({"sql", "name"})


def _sanitize_sql_lab_url_for_llm_context(url: str) -> str:
Expand Down Expand Up @@ -128,7 +128,10 @@ def open_sql_lab_with_context(
params["sql"] = request.sql

if request.title:
params["title"] = request.title
# SQL Lab's PopEditorTab reads `name` from query string to set the
# tab label; keep the MCP-facing field as `title` (more natural for
# LLMs) but emit `name` on the URL.
params["name"] = request.title
Comment thread
msyavuz marked this conversation as resolved.

if request.dataset_in_context:
# Add dataset context as a comment in the SQL if no SQL provided
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -143,9 +143,10 @@ def test_sanitizes_direct_sql_and_title_in_url_and_response(self) -> None:
assert parsed.path == "/sqllab"
assert params["dbid"] == ["7"]
assert params["schema"] == ["analytics"]
assert params["title"] == [
sanitize_for_llm_context("Review this query", field_path=("title",))
assert params["name"] == [
sanitize_for_llm_context("Review this query", field_path=("name",))
]
assert "title" not in params
assert params["sql"] == [
sanitize_for_llm_context(
"SELECT * FROM users LIMIT 10",
Expand Down Expand Up @@ -243,7 +244,7 @@ def test_sanitizes_sql_lab_url_query_parameters_for_llm_context(self) -> None:
try:
url = (
"https://superset.example.com/sqllab?"
"dbid=7&schema=analytics&sql=SELECT+1&title=Inspect+query"
"dbid=7&schema=analytics&sql=SELECT+1&name=Inspect+query"
)

response = mod._sanitize_sql_lab_response_for_llm_context(
Expand All @@ -261,8 +262,8 @@ def test_sanitizes_sql_lab_url_query_parameters_for_llm_context(self) -> None:
assert params["sql"] == [
sanitize_for_llm_context("SELECT 1", field_path=("sql",))
]
assert params["title"] == [
sanitize_for_llm_context("Inspect query", field_path=("title",))
assert params["name"] == [
sanitize_for_llm_context("Inspect query", field_path=("name",))
]
assert response.title == sanitize_for_llm_context(
"Inspect query",
Expand All @@ -271,6 +272,38 @@ def test_sanitizes_sql_lab_url_query_parameters_for_llm_context(self) -> None:
finally:
_restore_modules(saved_modules)

def test_whitespace_only_title_is_dropped(self) -> None:
"""Whitespace-only titles must not produce a blank-looking tab label."""
mod, saved_modules = _get_tool_module()
try:
request = OpenSqlLabRequest(
database_id=7,
sql="SELECT 1",
title=" ",
)

with (
patch(
"superset.daos.database.DatabaseDAO.find_by_id",
return_value=Mock(database_name="examples"),
),
patch.object(
mod.event_logger, "log_context", return_value=nullcontext()
),
patch.object(
mod,
"get_superset_base_url",
return_value="https://superset.example.com",
),
):
response = mod.open_sql_lab_with_context(request, _make_mock_ctx())

params = parse_qs(urlsplit(response.url).query)
assert "name" not in params
assert response.title is None
finally:
_restore_modules(saved_modules)

def test_sanitizes_error_and_keeps_empty_url_for_missing_database(self) -> None:
mod, saved_modules = _get_tool_module()
try:
Expand Down
Loading