Schema reads do not apply ROW_CAP and follow result chunks (#703) - #737
Merged
WaylandYang merged 1 commit intoSep 17, 2026
Conversation
…#703) A Databricks schema read on a wide catalog was cut at 201 rows silently. The same SQL Statement Execution API call goes through DatabricksEngine::run with row_limit = ROW_CAP + 1, intended as the chat-query safety valve. information_schema.columns returns one row per column, so any catalog with more than ~201 columns stops partway through a table and nothing says so. Snowflake has the same shape: the response includes partitionInfo with per-partition URLs that follow the first inline chunk, and the engine reads only that first chunk. Two fixes per engine, scoped to schema reads: - Schema reads do not send row_limit. fetch_schema calls run with None; chat-query execute() keeps Some(ROW_CAP + 1). Snowflake has no request-body row limit to begin with — it follows the existing LIMIT-clause path through wrap_limit — so this is a Databricks-only behavioural change. The row_limit is now an Option<usize> on run. - Both engines follow result chunks until exhausted. Databricks follows next_chunk_internal_link via repeated GET until the link is absent, then verifies rows.len() == manifest.total_row_count when no row_limit was set. Snowflake iterates partitionInfo from index 1 (the first partition is the inline data already in the response), GET each url, and verifies partition row counts sum to numRows. Any mismatch bails loudly instead of returning a partial schema. manifest.truncated on the Databricks side is now an explicit fail-fast: a chat query that actually gets truncated by the server returns an error rather than presenting fewer rows as a complete answer. The previous code returned the truncated set with no warning. Tests (all in wiremock, no live cluster needed): - a_schema_read_follows_chunks_until_exhausted — three chunks, all rows collected in order. - a_chat_query_bails_loudly_when_the_service_says_truncated — chat query with manifest.truncated=true errors instead of returning fewer rows. - a_schema_read_fails_when_total_row_count_does_not_match — schema read with a mismatch between manifest.total_row_count and the actual rows received errors. - a_schema_read_follows_partitions_until_exhausted — three partitions on the Snowflake side, all rows collected in order. Both new tests pass alongside the existing 9 (Databricks) and 3 (Snowflake) tests. utopia-server 304 passed, clippy -D warnings clean, cargo fmt clean. Integration test on a real Databricks / Snowflake warehouse is still required for end-to-end confirmation. Refs deeplethe#703 Signed-off-by: rollroyces <rollroyces@users.noreply.github.com>
rollroyces
force-pushed
the
fix/databricks-schema-row-limit-and-chunk-following
branch
2 times, most recently
from
September 16, 2026 15:37
2c52139 to
4ac7445
Compare
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
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.
Closes #703.
A Databricks schema read on a wide catalog was cut at 201 rows silently. The same SQL Statement Execution API call went through
DatabricksEngine::runwithrow_limit = ROW_CAP + 1, intended as the chat-query safety valve.information_schema.columnsreturns one row per column, so any catalog with more than ~201 columns stopped partway through a table and nothing said so. Snowflake has the same shape: the response includespartitionInfowith per-partition URLs following the first inline chunk, and the engine read only that first chunk.Fix
Two fixes per engine, scoped to schema reads:
row_limit.fetch_schemacallsrunwithNone;executekeepsSome(ROW_CAP + 1). The Databricksrunsignature is nowrun(sql: &str, row_limit: Option<usize>). Snowflake has no request-body row limit to begin with — the existingwrap_limitpath applies theLIMITclause toexecuteonly — so this is a Databricks-only behavioural change.next_chunk_internal_linkvia repeated GET until the link is absent, then verifiesrows.len() == manifest.total_row_countwhen norow_limitwas set. Snowflake iteratespartitionInfofrom index 1 (the first partition is the inline data already in the response), GETs eachurl, and verifies the partitionrowCounts sum tonumRows. Any mismatch bails loudly instead of returning a partial schema.manifest.truncatedon the Databricks side is now an explicit fail-fast: a chat query that actually gets truncated by the server returns an error rather than presenting fewer rows as a complete answer.Tests
All wiremock, no live cluster needed:
a_schema_read_follows_chunks_until_exhausted— three chunks, all rows collected in order.a_chat_query_bails_loudly_when_the_service_says_truncated— chat query withmanifest.truncated=trueerrors.a_schema_read_fails_when_total_row_count_does_not_match— schema read with a mismatch betweenmanifest.total_row_countand the actual rows errors.a_schema_read_follows_partitions_until_exhausted— three partitions on the Snowflake side, all rows collected in order.utopia-server304 passed (was 295 before — 4 new tests added);clippy -D warningsclean;cargo fmtclean.What still needs a real cluster
The two
fetch_schemaintegration tests still need a live Databricks and Snowflake warehouse, per #241 and #242. Those should keep the existing structure: run the engine against a real catalog with >201 columns (Databricks) and >1 partition (Snowflake), confirm the schema document reports every column.