fix(databricks): apply user-selected row limit to SQL Editor queries (#20190)

h3n4l · claude · web-flow · commit 4af45e510f5d · 2026-04-30T09:21:40.000Z
The Databricks driver discarded `db.QueryContext`, so the row-limit
toggle in the SQL Editor (default 1000) was silently ignored. Wide
queries against shared catalogs hit Databricks' INLINE response cap
(~25 MiB) and aborted with no result available:

    Statements executed with disposition=INLINE can have a result size
    of at most 20214400 bytes.

Forward `queryContext.Limit` into `execStatementSync` and set
`RowLimit` on the SDK request when non-zero. Sync and dump paths pass
0 so admin queries (SHOW CREATE TABLE, DESC FORMATTED, SELECT VERSION)
remain unbounded.

This also surfaces a `truncated` flag on the response when the limit
trims the result set.

Refs BYT-9379.

Co-authored-by: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/backend/plugin/db/databricks/databricks.go b/backend/plugin/db/databricks/databricks.go
@@ -78,7 +78,7 @@ func (*Driver) GetDB() *sql.DB {
 	return nil
 }
 
-func (d *Driver) QueryConn(ctx context.Context, _ *sql.Conn, statement string, _ db.QueryContext) ([]*v1pb.QueryResult, error) {
+func (d *Driver) QueryConn(ctx context.Context, _ *sql.Conn, statement string, queryContext db.QueryContext) ([]*v1pb.QueryResult, error) {
 	var results []*v1pb.QueryResult
 	stmts, err := base.SplitMultiSQL(storepb.Engine_DATABRICKS, statement)
 	if err != nil {
@@ -88,7 +88,7 @@ func (d *Driver) QueryConn(ctx context.Context, _ *sql.Conn, statement string, _
 	for _, stmt := range stmts {
 		result := &v1pb.QueryResult{}
 		startTime := time.Now()
-		dataArr, colInfo, err := d.execStatementSync(ctx, stmt.Text)
+		dataArr, colInfo, err := d.execStatementSync(ctx, stmt.Text, int64(queryContext.Limit))
 		if err != nil {
 			return nil, err
 		}
@@ -158,15 +158,25 @@ func (d *Driver) Execute(ctx context.Context, statement string, _ db.ExecuteOpti
 }
 
 // Execute SQL statement synchronously and return row data or error.
-func (d *Driver) execStatementSync(ctx context.Context, statement string) ([][]string, []dbsql.ColumnInfo, error) {
+// rowLimit caps the result set when > 0; 0 means no limit (used by sync
+// and dump paths). DDL/DML callers should pass 0.
+func (d *Driver) execStatementSync(ctx context.Context, statement string, rowLimit int64) ([][]string, []dbsql.ColumnInfo, error) {
 	// Bind the connected Bytebase database (= Unity Catalog) to the
 	// session, so unqualified `schema.table` references resolve like in
 	// other RDBMSes. Equivalent to issuing `USE CATALOG <name>` first.
-	resp, err := d.Client.StatementExecution.ExecuteAndWait(ctx, dbsql.ExecuteStatementRequest{
+	req := dbsql.ExecuteStatementRequest{
 		Statement:   statement,
 		WarehouseId: d.WarehouseID,
 		Catalog:     d.curCatalog,
-	})
+	}
+	// Honor the user-selected row limit (default 1000 from the SQL Editor).
+	// Without it, INLINE results above ~25 MiB are aborted server-side
+	// with no result available — see DispositionExternalLinks for the
+	// large-result path.
+	if rowLimit > 0 {
+		req.RowLimit = rowLimit
+	}
+	resp, err := d.Client.StatementExecution.ExecuteAndWait(ctx, req)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/backend/plugin/db/databricks/dump.go b/backend/plugin/db/databricks/dump.go
@@ -129,7 +129,7 @@ func (d *Driver) Dump(ctx context.Context, writer io.Writer, _ *storepb.Database
 }
 
 func (d *Driver) showCreateTable(ctx context.Context, qualifiedTblName string) (string, error) {
-	rows, colInfo, err := d.execStatementSync(ctx, fmt.Sprintf("SHOW CREATE TABLE %s", qualifiedTblName))
+	rows, colInfo, err := d.execStatementSync(ctx, fmt.Sprintf("SHOW CREATE TABLE %s", qualifiedTblName), 0)
 	if err != nil {
 		return "", err
 	}
@@ -170,7 +170,7 @@ func (d *Driver) descTbl(ctx context.Context, tblName string) (*tableAdditionalI
 	tblAddInfo := &tableAdditionalInfo{
 		tblProps: make(map[string]string),
 	}
-	rows, _, err := d.execStatementSync(ctx, fmt.Sprintf("DESC FORMATTED %s", tblName))
+	rows, _, err := d.execStatementSync(ctx, fmt.Sprintf("DESC FORMATTED %s", tblName), 0)
 	if err != nil {
 		return nil, err
 	}
diff --git a/backend/plugin/db/databricks/sync.go b/backend/plugin/db/databricks/sync.go
@@ -54,7 +54,7 @@ func (d *Driver) SyncInstance(ctx context.Context) (*db.InstanceMetadata, error)
 	instanceMetadata := &db.InstanceMetadata{}
 
 	// fetch version.
-	versionData, _, err := d.execStatementSync(ctx, "SELECT VERSION()")
+	versionData, _, err := d.execStatementSync(ctx, "SELECT VERSION()", 0)
 	if err != nil {
 		return nil, err
 	}

Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ func (d Driver) Dump(ctx context.Context, writer io.Writer, _ storepb.Database`
`129`	`129`	`}`
`130`	`130`
`131`	`131`	`func (d *Driver) showCreateTable(ctx context.Context, qualifiedTblName string) (string, error) {`
`132`		`- rows, colInfo, err := d.execStatementSync(ctx, fmt.Sprintf("SHOW CREATE TABLE %s", qualifiedTblName))`
	`132`	`+ rows, colInfo, err := d.execStatementSync(ctx, fmt.Sprintf("SHOW CREATE TABLE %s", qualifiedTblName), 0)`
`133`	`133`	`if err != nil {`
`134`	`134`	`return "", err`
`135`	`135`	`}`
`@@ -170,7 +170,7 @@ func (d Driver) descTbl(ctx context.Context, tblName string) (tableAdditionalI`
`170`	`170`	`tblAddInfo := &tableAdditionalInfo{`
`171`	`171`	`tblProps: make(map[string]string),`
`172`	`172`	`}`
`173`		`- rows, _, err := d.execStatementSync(ctx, fmt.Sprintf("DESC FORMATTED %s", tblName))`
	`173`	`+ rows, _, err := d.execStatementSync(ctx, fmt.Sprintf("DESC FORMATTED %s", tblName), 0)`
`174`	`174`	`if err != nil {`
`175`	`175`	`return nil, err`
`176`	`176`	`}`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ func (d Driver) SyncInstance(ctx context.Context) (db.InstanceMetadata, error)`
`54`	`54`	`instanceMetadata := &db.InstanceMetadata{}`
`55`	`55`
`56`	`56`	`// fetch version.`
`57`		`- versionData, _, err := d.execStatementSync(ctx, "SELECT VERSION()")`
	`57`	`+ versionData, _, err := d.execStatementSync(ctx, "SELECT VERSION()", 0)`
`58`	`58`	`if err != nil {`
`59`	`59`	`return nil, err`
`60`	`60`	`}`