fix: query cancellation leads to null results being cached (#6927)

k-anshul · k-anshul · commit 9ae132dc211f · 2025-03-19T19:27:24.000+05:30
* fix: query cancellation leads to null results being cached

* fix unit tests and lint failures

* fix some more instances
diff --git a/.golangci.yml b/.golangci.yml
@@ -47,6 +47,7 @@ linters:
     - unused
     - usestdlibvars
     - whitespace
+    - rowserrcheck
 
 linters-settings:
   depguard:
diff --git a/admin/provisioner/clickhousestatic/provisioner_test.go b/admin/provisioner/clickhousestatic/provisioner_test.go
@@ -74,6 +74,7 @@ func TestClickHouseStatic(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, count, 1)
 	}
+	require.NoError(t, err)
 	rows.Close()
 
 	// Get the name of the first connection's database
@@ -93,6 +94,7 @@ func TestClickHouseStatic(t *testing.T) {
 	for rows.Next() {
 		require.Fail(t, "unexpected visible table in information schema")
 	}
+	require.NoError(t, err)
 	rows.Close()
 
 	// Deprovision the resources
diff --git a/runtime/drivers/clickhouse/olap.go b/runtime/drivers/clickhouse/olap.go
@@ -479,6 +479,10 @@ func (c *connection) RenameTable(ctx context.Context, oldName, newName string) e
 				return err
 			}
 		}
+		err = res.Err()
+		if err != nil {
+			return err
+		}
 		res.Close()
 		engineFull = strings.ReplaceAll(engineFull, localTableName(oldName), safelocalTableName(newName))
 
@@ -535,6 +539,10 @@ func (c *connection) renameView(ctx context.Context, oldName, newName, onCluster
 			return err
 		}
 	}
+	err = res.Err()
+	if err != nil {
+		return err
+	}
 	res.Close()
 
 	// create new view
@@ -735,6 +743,10 @@ func (c *connection) columnClause(ctx context.Context, table string) (string, er
 		columnClause.WriteString(" ")
 		columnClause.WriteString(typ)
 	}
+	err = res.Err()
+	if err != nil {
+		return "", err
+	}
 	columnClause.WriteRune(')')
 	return columnClause.String(), nil
 }
@@ -838,6 +850,10 @@ func (c *connection) getTableEngine(ctx context.Context, name string) (string, e
 			return "", err
 		}
 	}
+	err = res.Err()
+	if err != nil {
+		return "", err
+	}
 	return engine, nil
 }
 
@@ -860,6 +876,10 @@ func (c *connection) getTablePartitions(ctx context.Context, name string) ([]str
 		}
 		partitions = append(partitions, part)
 	}
+	err = res.Err()
+	if err != nil {
+		return nil, err
+	}
 	return partitions, nil
 }
 
diff --git a/runtime/drivers/clickhouse/olap_test.go b/runtime/drivers/clickhouse/olap_test.go
@@ -194,6 +194,7 @@ func testInsertTableAsSelect_WithAppend(t *testing.T, olap drivers.OLAPStore) {
 		require.NoError(t, res.Scan(&r.ID, &r.Planet))
 		result = append(result, r)
 	}
+	require.NoError(t, err)
 
 	expected := []struct {
 		ID     int
@@ -258,6 +259,7 @@ func testInsertTableAsSelect_WithMerge(t *testing.T, olap drivers.OLAPStore) {
 		require.NoError(t, res.Scan(&r.ID, &r.Value))
 		result = append(result, r)
 	}
+	require.NoError(t, err)
 
 	expected := map[int]string{
 		0: "insert",
@@ -328,6 +330,7 @@ func testInsertTableAsSelect_WithPartitionOverwrite(t *testing.T, olap drivers.O
 		require.NoError(t, res.Scan(&r.ID, &r.Value))
 		result = append(result, r)
 	}
+	require.NoError(t, err)
 
 	expected := []struct {
 		ID    int
@@ -394,6 +397,7 @@ func testInsertTableAsSelect_WithPartitionOverwrite_DatePartition(t *testing.T,
 		require.NoError(t, res.Scan(&r.DT, &r.Value))
 		result = append(result, r)
 	}
+	require.NoError(t, err)
 
 	expected := []struct {
 		DT    string
diff --git a/runtime/drivers/druid/druid_container_test.go b/runtime/drivers/druid/druid_container_test.go
@@ -179,6 +179,7 @@ func testTimeFloor(t *testing.T, olap drivers.OLAPStore) {
 		require.Equal(t, 15, tm.Minute())
 		count += 1
 	}
+	require.NoError(t, rows.Err())
 	require.Equal(t, 9, count)
 }
 
diff --git a/runtime/drivers/duckdb/model_executor_https_self_test.go b/runtime/drivers/duckdb/model_executor_https_self_test.go
@@ -100,5 +100,6 @@ func TestHTTPToDuckDBTransfer(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, 3, count)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 }
diff --git a/runtime/drivers/duckdb/olap_crud_test.go b/runtime/drivers/duckdb/olap_crud_test.go
@@ -198,6 +198,7 @@ func Test_connection_InsertTableAsSelect_WithMergeStrategy(t *testing.T) {
 		require.NoError(t, res.Scan(&r.Range, &r.Strategy))
 		results = append(results, r)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 
 	exptected := []struct {
diff --git a/runtime/drivers/duckdb/transporter_mysql_to_duckDB_test.go b/runtime/drivers/duckdb/transporter_mysql_to_duckDB_test.go
@@ -129,6 +129,7 @@ func allMySQLDataTypesTest(t *testing.T, db *sql.DB, dsn string) {
 		require.NoError(t, err)
 		require.Equal(t, count, 2)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 	require.NoError(t, to.Close())
 }
@@ -180,6 +181,7 @@ func mysqlToDuckDB(t *testing.T, dsn string) {
 		require.NoError(t, err)
 		require.Equal(t, 2, count)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 	tbl, err := olap.InformationSchema().Lookup(context.Background(), "", "", "sink")
 	require.NoError(t, err)
diff --git a/runtime/drivers/duckdb/transporter_postgres_to_duckDB_test.go b/runtime/drivers/duckdb/transporter_postgres_to_duckDB_test.go
@@ -90,6 +90,7 @@ func allDataTypesTest(t *testing.T, db *sql.DB, dbURL string) {
 		require.NoError(t, err)
 		require.Equal(t, 1, count)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 	require.NoError(t, to.Close())
 }
@@ -141,6 +142,7 @@ func pgxToDuckDB(t *testing.T, pgdb *sql.DB, dbURL string) {
 		require.NoError(t, err)
 		require.Equal(t, 1, count)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 
 	// ingest some more data in postges
@@ -165,6 +167,7 @@ func pgxToDuckDB(t *testing.T, pgdb *sql.DB, dbURL string) {
 		require.NoError(t, err)
 		require.Equal(t, 2, count)
 	}
+	require.NoError(t, res.Err())
 	require.NoError(t, res.Close())
 
 	require.NoError(t, duckDB.Close())
diff --git a/runtime/drivers/duckdb/transporter_test.go b/runtime/drivers/duckdb/transporter_test.go
@@ -195,6 +195,7 @@ mum,8.2`)
 			for rows.Next() {
 				colCount++
 			}
+			require.NoError(t, rows.Err())
 			require.NoError(t, rows.Close())
 			require.Equal(t, test.colCount, colCount)
 		})
@@ -438,6 +439,7 @@ func TestIterativeParquetIngestionWithVariableSchema(t *testing.T) {
 			for rows.Next() {
 				colCount++
 			}
+			require.NoError(t, rows.Err())
 			require.NoError(t, rows.Close())
 			require.Equal(t, test.colCount, colCount)
 		})
@@ -586,6 +588,7 @@ func TestIterativeJSONIngestionWithVariableSchema(t *testing.T) {
 			for rows.Next() {
 				colCount++
 			}
+			require.NoError(t, rows.Err())
 			require.NoError(t, rows.Close())
 			require.Equal(t, test.colCount, colCount)
 		})
diff --git a/runtime/drivers/olap.go b/runtime/drivers/olap.go
@@ -641,6 +641,10 @@ func (d Dialect) SelectInlineResults(result *Result) (string, []any, []any, erro
 
 		rows++
 	}
+	err := result.Err()
+	if err != nil {
+		return "", nil, nil, err
+	}
 
 	if d == DialectDruid || d == DialectDuckDB || d == DialectPinot {
 		prefix += ") "
diff --git a/runtime/drivers/sqlite/registry.go b/runtime/drivers/sqlite/registry.go
@@ -129,6 +129,9 @@ func (c *connection) findInstances(_ context.Context, whereClause string, args .
 
 		res = append(res, i)
 	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
 
 	return res, nil
 }
diff --git a/runtime/metricsview/executor.go b/runtime/metricsview/executor.go
@@ -117,7 +117,7 @@ func (e *Executor) CacheKey(ctx context.Context) ([]byte, bool, error) {
 			return nil, false, err
 		}
 	}
-	if res.Err() != nil {
+	if err := res.Err(); err != nil {
 		return nil, false, err
 	}
 
@@ -523,8 +523,9 @@ func (e *Executor) Search(ctx context.Context, qry *SearchQuery, executionTime *
 		}
 		searchResult = append(searchResult, row)
 	}
-	if res.Err() != nil {
-		return nil, res.Err()
+	err = res.Err()
+	if err != nil {
+		return nil, err
 	}
 	return searchResult, nil
 }
diff --git a/runtime/metricsview/executor_rewrite_druid_exactify.go b/runtime/metricsview/executor_rewrite_druid_exactify.go
@@ -99,6 +99,10 @@ func (e *Executor) rewriteQueryDruidExactify(ctx context.Context, qry *Query) er
 
 		vals = append(vals, val)
 	}
+	err = res.Err()
+	if err != nil {
+		return err
+	}
 
 	// Add the dimensions values as a "<dim> IN (<vals...>)" expression in the outer query's WHERE clause.
 	var inExpr *Expression
diff --git a/runtime/pkg/duckdbsql/duckdbsql.go b/runtime/pkg/duckdbsql/duckdbsql.go
@@ -20,12 +20,15 @@ func queryString(qry string, args ...any) ([]byte, error) {
 	defer func() { _ = rows.Close() }()
 
 	var res []byte
-	if rows.Next() {
+	for rows.Next() {
 		err := rows.Scan(&res)
 		if err != nil {
 			return nil, err
 		}
 	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
 
 	return res, nil
 }
diff --git a/runtime/pkg/duckdbsql/expression.go b/runtime/pkg/duckdbsql/expression.go
@@ -33,12 +33,15 @@ func queryBoolPair(qry string, args ...any) (bool, bool, error) {
 	defer func() { _ = rows.Close() }()
 
 	var a, b bool
-	if rows.Next() {
+	for rows.Next() {
 		err := rows.Scan(&a, &b)
 		if err != nil {
 			return false, false, err
 		}
 	}
+	if err := rows.Err(); err != nil {
+		return false, false, err
+	}
 
 	return a, b, nil
 }
diff --git a/runtime/pkg/fieldselectorpb/duckdb.go b/runtime/pkg/fieldselectorpb/duckdb.go
@@ -47,7 +47,7 @@ func resolveDuckDBExpression(expr string, all []string) ([]string, error) {
 			return err
 		}
 
-		return nil
+		return rows.Err()
 	})
 	if err != nil {
 		return nil, err
diff --git a/runtime/queries/column_numeric_histogram.go b/runtime/queries/column_numeric_histogram.go
@@ -433,6 +433,10 @@ func getMinMaxRange(ctx context.Context, olap drivers.OLAPStore, columnName, dat
 			return nil, nil, nil, err
 		}
 	}
+	err = minMaxRow.Err()
+	if err != nil {
+		return nil, nil, nil, err
+	}
 
 	minMaxRow.Close()
 
diff --git a/runtime/queries/column_timeseries.go b/runtime/queries/column_timeseries.go
@@ -500,6 +500,10 @@ func (q *ColumnTimeseries) CreateTimestampRollupReduction(
 
 			results = append(results, tsv)
 		}
+		err = rows.Err()
+		if err != nil {
+			return nil, err
+		}
 
 		return results, nil
 	}
@@ -592,6 +596,10 @@ func (q *ColumnTimeseries) CreateTimestampRollupReduction(
 			results[i-3], results[i-2] = results[i-2], results[i-3]
 		}
 	}
+	err = rows.Err()
+	if err != nil {
+		return nil, err
+	}
 
 	return results, nil
 }
diff --git a/runtime/queries/metricsview_timeseries.go b/runtime/queries/metricsview_timeseries.go
@@ -234,6 +234,10 @@ func (q *MetricsViewTimeSeries) populateResult(rows *drivers.Result, tsAlias str
 		})
 		start = addTo(t, dur, tz)
 	}
+	err := rows.Err()
+	if err != nil {
+		return err
+	}
 	if q.TimeEnd != nil && nullRecords != nil {
 		if start.Equal(zeroTime) && q.TimeStart != nil {
 			start = q.TimeStart.AsTime()
diff --git a/runtime/queries/resource_watermark.go b/runtime/queries/resource_watermark.go
@@ -112,6 +112,10 @@ func (q *ResourceWatermark) resolveMetricsView(ctx context.Context, rt *runtime.
 			return err
 		}
 	}
+	err = res.Err()
+	if err != nil {
+		return err
+	}
 
 	if !t.IsZero() {
 		// Hacky workaround for the following issue: the watermark is used as the *exclusive* upper bound for time ranges.
diff --git a/runtime/queries/table_columns.go b/runtime/queries/table_columns.go
@@ -117,6 +117,10 @@ func (q *TableColumns) Resolve(ctx context.Context, rt *runtime.Runtime, instanc
 				pcs = append(pcs, &pc)
 				i++
 			}
+			err = rows.Err()
+			if err != nil {
+				return err
+			}
 
 			q.Result = &runtimev1.TableColumnsResponse{
 				ProfileColumns: pcs[0:i],
diff --git a/runtime/server/queries_columns_test.go b/runtime/server/queries_columns_test.go
@@ -615,6 +615,7 @@ func getColumnTestServerWithModel(t *testing.T, sql string, expectation int) (*s
 		err := res.Scan(&n)
 		require.NoError(t, err)
 	}
+	require.NoError(t, res.Err())
 	if expectation >= 0 {
 		require.Equal(t, expectation, n)
 	}
diff --git a/runtime/server/queries_tables_test.go b/runtime/server/queries_tables_test.go
@@ -7,7 +7,6 @@ import (
 	"github.com/rilldata/rill/runtime/pkg/ratelimit"
 	"github.com/rilldata/rill/runtime/server"
 
-	"github.com/jmoiron/sqlx"
 	runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
 	"github.com/rilldata/rill/runtime/pkg/activity"
 	"github.com/rilldata/rill/runtime/testruntime"
@@ -132,16 +131,6 @@ func getTableTestServerWithSql(t *testing.T, sql string) (*server.Server, string
 	return server, instanceID
 }
 
-func getSingleValue(t *testing.T, rows *sqlx.Rows) int {
-	var val int
-	if rows.Next() {
-		err := rows.Scan(&val)
-		require.NoError(t, err)
-	}
-	rows.Close()
-	return val
-}
-
 func getTableTestServerWithEmptyModel(t *testing.T) (*server.Server, string) {
 	rt, instanceID := testruntime.NewInstanceWithModel(t, "test", `
 		SELECT 1::int AS a, 10::int AS "b""b" where 1<>1
diff --git a/runtime/testruntime/olap.go b/runtime/testruntime/olap.go

Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,7 @@ func TestClickHouseStatic(t *testing.T) {`
`74`	`74`	`require.NoError(t, err)`
`75`	`75`	`require.Equal(t, count, 1)`
`76`	`76`	`}`
	`77`	`+ require.NoError(t, err)`
`77`	`78`	`rows.Close()`
`78`	`79`
`79`	`80`	`// Get the name of the first connection's database`
`@@ -93,6 +94,7 @@ func TestClickHouseStatic(t *testing.T) {`
`93`	`94`	`for rows.Next() {`
`94`	`95`	`require.Fail(t, "unexpected visible table in information schema")`
`95`	`96`	`}`
	`97`	`+ require.NoError(t, err)`
`96`	`98`	`rows.Close()`
`97`	`99`
`98`	`100`	`// Deprovision the resources`
Original file line number	Diff line number	Diff line change
`@@ -479,6 +479,10 @@ func (c *connection) RenameTable(ctx context.Context, oldName, newName string) e`
`479`	`479`	`return err`
`480`	`480`	`}`
`481`	`481`	`}`
	`482`	`+ err = res.Err()`
	`483`	`+ if err != nil {`
	`484`	`+ return err`
	`485`	`+ }`
`482`	`486`	`res.Close()`
`483`	`487`	`engineFull = strings.ReplaceAll(engineFull, localTableName(oldName), safelocalTableName(newName))`
`484`	`488`
`@@ -535,6 +539,10 @@ func (c *connection) renameView(ctx context.Context, oldName, newName, onCluster`
`535`	`539`	`return err`
`536`	`540`	`}`
`537`	`541`	`}`
	`542`	`+ err = res.Err()`
	`543`	`+ if err != nil {`
	`544`	`+ return err`
	`545`	`+ }`
`538`	`546`	`res.Close()`
`539`	`547`
`540`	`548`	`// create new view`
`@@ -735,6 +743,10 @@ func (c *connection) columnClause(ctx context.Context, table string) (string, er`
`735`	`743`	`columnClause.WriteString(" ")`
`736`	`744`	`columnClause.WriteString(typ)`
`737`	`745`	`}`
	`746`	`+ err = res.Err()`
	`747`	`+ if err != nil {`
	`748`	`+ return "", err`
	`749`	`+ }`
`738`	`750`	`columnClause.WriteRune(')')`
`739`	`751`	`return columnClause.String(), nil`
`740`	`752`	`}`
`@@ -838,6 +850,10 @@ func (c *connection) getTableEngine(ctx context.Context, name string) (string, e`
`838`	`850`	`return "", err`
`839`	`851`	`}`
`840`	`852`	`}`
	`853`	`+ err = res.Err()`
	`854`	`+ if err != nil {`
	`855`	`+ return "", err`
	`856`	`+ }`
`841`	`857`	`return engine, nil`
`842`	`858`	`}`
`843`	`859`
`@@ -860,6 +876,10 @@ func (c *connection) getTablePartitions(ctx context.Context, name string) ([]str`
`860`	`876`	`}`
`861`	`877`	`partitions = append(partitions, part)`
`862`	`878`	`}`
	`879`	`+ err = res.Err()`
	`880`	`+ if err != nil {`
	`881`	`+ return nil, err`
	`882`	`+ }`
`863`	`883`	`return partitions, nil`
`864`	`884`	`}`
`865`	`885`
Original file line number	Diff line number	Diff line change
`@@ -179,6 +179,7 @@ func testTimeFloor(t *testing.T, olap drivers.OLAPStore) {`
`179`	`179`	`require.Equal(t, 15, tm.Minute())`
`180`	`180`	`count += 1`
`181`	`181`	`}`
	`182`	`+ require.NoError(t, rows.Err())`
`182`	`183`	`require.Equal(t, 9, count)`
`183`	`184`	`}`
`184`	`185`
Original file line number	Diff line number	Diff line change
`@@ -100,5 +100,6 @@ func TestHTTPToDuckDBTransfer(t *testing.T) {`
`100`	`100`	`require.NoError(t, err)`
`101`	`101`	`require.Equal(t, 3, count)`
`102`	`102`	`}`
	`103`	`+ require.NoError(t, res.Err())`
`103`	`104`	`require.NoError(t, res.Close())`
`104`	`105`	`}`
Original file line number	Diff line number	Diff line change
`@@ -198,6 +198,7 @@ func Test_connection_InsertTableAsSelect_WithMergeStrategy(t *testing.T) {`
`198`	`198`	`require.NoError(t, res.Scan(&r.Range, &r.Strategy))`
`199`	`199`	`results = append(results, r)`
`200`	`200`	`}`
	`201`	`+ require.NoError(t, res.Err())`
`201`	`202`	`require.NoError(t, res.Close())`
`202`	`203`
`203`	`204`	`exptected := []struct {`
Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,7 @@ func allMySQLDataTypesTest(t testing.T, db sql.DB, dsn string) {`
`129`	`129`	`require.NoError(t, err)`
`130`	`130`	`require.Equal(t, count, 2)`
`131`	`131`	`}`
	`132`	`+ require.NoError(t, res.Err())`
`132`	`133`	`require.NoError(t, res.Close())`
`133`	`134`	`require.NoError(t, to.Close())`
`134`	`135`	`}`
`@@ -180,6 +181,7 @@ func mysqlToDuckDB(t *testing.T, dsn string) {`
`180`	`181`	`require.NoError(t, err)`
`181`	`182`	`require.Equal(t, 2, count)`
`182`	`183`	`}`
	`184`	`+ require.NoError(t, res.Err())`
`183`	`185`	`require.NoError(t, res.Close())`
`184`	`186`	`tbl, err := olap.InformationSchema().Lookup(context.Background(), "", "", "sink")`
`185`	`187`	`require.NoError(t, err)`
Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,7 @@ func allDataTypesTest(t testing.T, db sql.DB, dbURL string) {`
`90`	`90`	`require.NoError(t, err)`
`91`	`91`	`require.Equal(t, 1, count)`
`92`	`92`	`}`
	`93`	`+ require.NoError(t, res.Err())`
`93`	`94`	`require.NoError(t, res.Close())`
`94`	`95`	`require.NoError(t, to.Close())`
`95`	`96`	`}`
`@@ -141,6 +142,7 @@ func pgxToDuckDB(t testing.T, pgdb sql.DB, dbURL string) {`
`141`	`142`	`require.NoError(t, err)`
`142`	`143`	`require.Equal(t, 1, count)`
`143`	`144`	`}`
	`145`	`+ require.NoError(t, res.Err())`
`144`	`146`	`require.NoError(t, res.Close())`
`145`	`147`
`146`	`148`	`// ingest some more data in postges`
`@@ -165,6 +167,7 @@ func pgxToDuckDB(t testing.T, pgdb sql.DB, dbURL string) {`
`165`	`167`	`require.NoError(t, err)`
`166`	`168`	`require.Equal(t, 2, count)`
`167`	`169`	`}`
	`170`	`+ require.NoError(t, res.Err())`
`168`	`171`	`require.NoError(t, res.Close())`
`169`	`172`
`170`	`173`	`require.NoError(t, duckDB.Close())`