fix(mcp): restore connectionId on execute-sql tool#60454
Merged
Conversation
Restores the ability to target an external data warehouse connection (e.g. DuckLake, direct Postgres) from the MCP that was lost when query-run was replaced by execute-sql. Adds an optional connectionId arg to ExecuteSQLMCPToolArgs and the matching Zod schema. When set, the tool wraps the query in a HogQLQuery(connectionId=...) and bypasses the local HogQL parse/print validation step — those queries reference tables that don't exist in the default ClickHouse database, so validation must be deferred to HogQLQueryRunner, which already resolves schema for the selected connection. When connectionId is omitted, behavior is unchanged. Generated-By: PostHog Code Task-Id: e582d64f-f10d-4f6f-a4cb-2656def3b12c
Contributor
Query snapshots: Backend query snapshots updatedChanges: 1 snapshots (1 modified, 0 added, 0 deleted) What this means:
Next steps:
|
Contributor
|
Size Change: 0 B Total Size: 80.6 MB ℹ️ View Unchanged
|
Regenerated services/mcp/schema/tool-inputs.json from the updated Zod schemas via `services/mcp/scripts/generate-tool-schema.ts` so the generated artifact matches the new optional `connectionId` field on `ExecuteSQLSchema`. CI's "Check MCP schema is up to date" step flagged the drift. Generated-By: PostHog Code Task-Id: e582d64f-f10d-4f6f-a4cb-2656def3b12c
Mirrors the wording used by `query-validate` and `hogql-schema`, the other flat-query HogQL tools that also expose a top-level `connectionId` field. Most importantly, adds the "Use external-data-sources-list to discover available connection ids" hint so agents know how to find a valid id. Updates the Zod schema, the regenerated tool-inputs.json, the Python Pydantic args, and the v2 snapshot to match. Generated-By: PostHog Code Task-Id: e582d64f-f10d-4f6f-a4cb-2656def3b12c
georgemunyoro
approved these changes
May 28, 2026
Contributor
Prompt To Fix All With AIFix the following 2 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 2
ee/hogai/tools/execute_sql/test/test_mcp_tool.py:93-97
The custom instruction says to always prefer parameterised tests. This test covers only `" "` (spaces), but the empty-check code path uses `rstrip(";").strip()` — other blank-like inputs such as `""`, `";"`, `" ; "` are all worth verifying. Note that `" ; "` specifically would *not* be caught: `" ; ".rstrip(";")` leaves `" ; "` (trailing char is space), then `.strip()` yields `";"`, which is truthy — so a query of `";"` surrounded by spaces is silently forwarded to the runner. A parameterised test would surface this.
```suggestion
@parameterized.expand(
[
("empty_string", ""),
("spaces_only", " "),
("semicolon_only", ";"),
("semicolon_with_leading_spaces", " ;"),
]
)
async def test_connection_id_with_empty_query_raises(self, _name: str, query: str):
with self.assertRaises(MaxToolRetryableError):
await self.tool.execute(
ExecuteSQLMCPToolArgs(query=query, connectionId="conn_abc"),
)
```
### Issue 2 of 2
ee/hogai/tools/execute_sql/mcp_tool.py:46
The `rstrip(";").strip()` order has a subtle gap: a query like `" ; "` (semicolon with trailing spaces) produces `";"` after this chain — `rstrip(";")` removes nothing because the last character is a space, then `strip()` yields the non-empty `";"`. The fast-fail empty-query guard is bypassed and the runner receives `";"` as the query, which then fails with a less clear HogQL parse error instead of the expected `"Query is empty"` message. Reversing the order to `strip().rstrip(";").strip()` catches this case.
```suggestion
cleaned_query = args.query.strip().rstrip(";").strip() if args.query else ""
```
Reviews (1): Last reviewed commit: "chore(mcp): align connectionId descripti..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+93
to
+97
| async def test_connection_id_with_empty_query_raises(self): | ||
| with self.assertRaises(MaxToolRetryableError): | ||
| await self.tool.execute( | ||
| ExecuteSQLMCPToolArgs(query=" ", connectionId="conn_abc"), | ||
| ) |
Contributor
There was a problem hiding this comment.
" " (spaces), but the empty-check code path uses rstrip(";").strip() — other blank-like inputs such as "", ";", " ; " are all worth verifying. Note that " ; " specifically would not be caught: " ; ".rstrip(";") leaves " ; " (trailing char is space), then .strip() yields ";", which is truthy — so a query of ";" surrounded by spaces is silently forwarded to the runner. A parameterised test would surface this.
Suggested change
| async def test_connection_id_with_empty_query_raises(self): | |
| with self.assertRaises(MaxToolRetryableError): | |
| await self.tool.execute( | |
| ExecuteSQLMCPToolArgs(query=" ", connectionId="conn_abc"), | |
| ) | |
| @parameterized.expand( | |
| [ | |
| ("empty_string", ""), | |
| ("spaces_only", " "), | |
| ("semicolon_only", ";"), | |
| ("semicolon_with_leading_spaces", " ;"), | |
| ] | |
| ) | |
| async def test_connection_id_with_empty_query_raises(self, _name: str, query: str): | |
| with self.assertRaises(MaxToolRetryableError): | |
| await self.tool.execute( | |
| ExecuteSQLMCPToolArgs(query=query, connectionId="conn_abc"), | |
| ) |
Prompt To Fix With AI
This is a comment left during a code review.
Path: ee/hogai/tools/execute_sql/test/test_mcp_tool.py
Line: 93-97
Comment:
The custom instruction says to always prefer parameterised tests. This test covers only `" "` (spaces), but the empty-check code path uses `rstrip(";").strip()` — other blank-like inputs such as `""`, `";"`, `" ; "` are all worth verifying. Note that `" ; "` specifically would *not* be caught: `" ; ".rstrip(";")` leaves `" ; "` (trailing char is space), then `.strip()` yields `";"`, which is truthy — so a query of `";"` surrounded by spaces is silently forwarded to the runner. A parameterised test would surface this.
```suggestion
@parameterized.expand(
[
("empty_string", ""),
("spaces_only", " "),
("semicolon_only", ";"),
("semicolon_with_leading_spaces", " ;"),
]
)
async def test_connection_id_with_empty_query_raises(self, _name: str, query: str):
with self.assertRaises(MaxToolRetryableError):
await self.tool.execute(
ExecuteSQLMCPToolArgs(query=query, connectionId="conn_abc"),
)
```
How can I resolve this? If you propose a fix, please make it concise.| # default ClickHouse database, so the local parse/print HogQL validation step | ||
| # would reject them. Defer validation to the runner, which resolves the schema | ||
| # for the selected connection. | ||
| cleaned_query = args.query.rstrip(";").strip() if args.query else "" |
Contributor
There was a problem hiding this comment.
rstrip(";").strip() order has a subtle gap: a query like " ; " (semicolon with trailing spaces) produces ";" after this chain — rstrip(";") removes nothing because the last character is a space, then strip() yields the non-empty ";". The fast-fail empty-query guard is bypassed and the runner receives ";" as the query, which then fails with a less clear HogQL parse error instead of the expected "Query is empty" message. Reversing the order to strip().rstrip(";").strip() catches this case.
Suggested change
| cleaned_query = args.query.rstrip(";").strip() if args.query else "" | |
| cleaned_query = args.query.strip().rstrip(";").strip() if args.query else "" |
Prompt To Fix With AI
This is a comment left during a code review.
Path: ee/hogai/tools/execute_sql/mcp_tool.py
Line: 46
Comment:
The `rstrip(";").strip()` order has a subtle gap: a query like `" ; "` (semicolon with trailing spaces) produces `";"` after this chain — `rstrip(";")` removes nothing because the last character is a space, then `strip()` yields the non-empty `";"`. The fast-fail empty-query guard is bypassed and the runner receives `";"` as the query, which then fails with a less clear HogQL parse error instead of the expected `"Query is empty"` message. Reversing the order to `strip().rstrip(";").strip()` catches this case.
```suggestion
cleaned_query = args.query.strip().rstrip(";").strip() if args.query else ""
```
How can I resolve this? If you propose a fix, please make it concise.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The
query-runMCP tool accepted aHogQLQuerybody, which carried an optionalconnectionIdfor routing queries against an external data warehouse connection (e.g. DuckLake, direct Postgres). The replacementexecute-sqltool's input schema is just{ query, truncate }— there is no way to target a non-default connection, so users that relied onconnectionIdcan no longer query their external warehouse via the MCP.Changes
ExecuteSQLMCPToolArgsgains an optionalconnectionId: str | Nonefield.connectionIdis set, the tool wraps the query in a realHogQLQuery(query=..., connectionId=...)and skips the local HogQL parse/print validation step. That validation builds the default ClickHouse database for table resolution, so queries against external-connection-only tables would always be rejected before reaching the runner.HogQLQueryRunneralready resolves the schema for the selected connection at execute time, so this is the right layer to validate.connectionIdis omitted, behaviour is unchanged — the existingAssistantHogQLQueryvalidation path runs.connectionIdfield added to the ZodExecuteSQLSchemaand forwarded through the TS handler.End-to-end flow with
connectionId: TS handler →invokeMcpToolPOSTs args → Pydantic validates →HogQLQuery(query, connectionId)is handed toInsightContext→execute_and_format_querydumps + re-validates, falling throughAssistantSupportedQueryRoot(which rejects the extraconnectionIdviaextra="forbid") toQuerySchemaRoot, which accepts the realHogQLQuery→HogQLQueryRunnerreadsself.query.connectionIdand routes viaget_direct_external_data_source_for_connection.How did you test this code?
I'm an agent (Claude / Claude Code) — no manual testing performed. I added two unit tests covering the new path:
test_connection_id_skips_local_validation_and_wraps_in_hogql_query— asserts that_validate_hogql_queryis not called whenconnectionIdis set, and that the query reachingInsightContextis aHogQLQuerycarrying the suppliedconnectionId.test_connection_id_with_empty_query_raises— asserts the empty-query guard still fires whenconnectionIdis set.I did not run the test suite, ruff, or
pnpm typescript:checklocally — flox/uv weren't available in the agent environment. CI is the source of truth for this PR.Publish to changelog?
no
Docs update
n/a
🤖 Agent context
Authored by Claude Code (PostHog Code task
e582d64f-f10d-4f6f-a4cb-2656def3b12c) in response to a Slack report from a user that they could no longer query their external warehouse via the MCP afterquery-runwas replaced byexecute-sql.Decisions along the way:
connectionId.HogQLQueryalso exposessendRawQuery, which is typically wanted for native Postgres passthrough on a direct connection. I considered adding it too but the original report was specifically aboutconnectionId; if users need raw-SQL passthrough as well, it's a small follow-up.connectionIdis set._validate_hogql_query_synccallsprepare_and_print_astagainst the default ClickHouse database, so any reference to a table that only exists on the external connection would always fail. The runner already handles schema resolution per-connection, so validation belongs there, not here. The empty-query guard is preserved on both branches.HogQLQuery(notAssistantHogQLQuery).AssistantHogQLQueryisextra="forbid"and has noconnectionIdfield, so it would strip the value. PassingHogQLQueryworks becauseexecute_and_format_queryre-validates withAssistantSupportedQueryRootfirst (which fails on the extra field) and then falls back toQuerySchemaRoot, which has the realHogQLQueryin its union and preservesconnectionId.connectionId(camelCase) to mirror the existingHogQLQuery.connectionIdfield rather than introducing a snake_case alias.