[v0.6] Prevent queries being interrupted by the sqlite layer (#897)

Author: aruiz14

pkg/sqlcache/db/client.go

@@ -208,7 +208,16 @@ func (c *client) QueryForRows(ctx context.Context, stmt transaction.Stmt, params
 	c.connLock.RLock()
 	defer c.connLock.RUnlock()
 
-	return stmt.QueryContext(ctx, params...)
+	// The underlying sqlite implementation seems to not cleanly close transactions when a query is interrupted
+	// We use context.Background to let the query finish correctly, immediately checking the original context afterward and properly closing Rows if canceled
+	rows, err := stmt.QueryContext(context.Background(), params...)
+	if err != nil {
+		return nil, err
+	} else if ctx.Err() != nil {
+		rows.Close()
+		return nil, ctx.Err()
+	}
+	return rows, nil
 }
 
 // CloseStmt will call close on the given Closable. It is intended to be used with a sql statement. This function is meant

pkg/sqlcache/db/transaction/transaction.go

@@ -39,7 +39,5 @@ func (c client) Stmt(stmt *sql.Stmt) Stmt {
 // rationale: allow mocking
 type Stmt interface {
 	Exec(args ...any) (sql.Result, error)
-	Query(args ...any) (*sql.Rows, error)
 	QueryContext(ctx context.Context, args ...any) (*sql.Rows, error)
-	QueryRowContext(ctx context.Context, args ...any) *sql.Row
 }

pkg/sqlcache/db/transaction_mocks_test.go

@@ -363,23 +363,32 @@ func (l *ListOptionIndexer) Watch(ctx context.Context, opts WatchOptions, events
 	}
 
 	if err := l.WithTransaction(ctx, false, func(tx transaction.Client) error {
-		rowIDRow := tx.Stmt(l.findEventsRowByRVStmt).QueryRowContext(ctx, targetRV)
-		if err := rowIDRow.Err(); err != nil {
-			return &db.QueryError{QueryString: l.findEventsRowByRVQuery, Err: err}
-		}
-
 		var rowID int
-		err := rowIDRow.Scan(&rowID)
-		if errors.Is(err, sql.ErrNoRows) {
-			if targetRV != latestRV {
-				return ErrTooOld
+		// use a closure to ensure rows is always closed immediately after it's needed
+		if err := func() error {
+			rows, err := l.QueryForRows(ctx, tx.Stmt(l.findEventsRowByRVStmt), targetRV)
+			if err != nil {
+				return &db.QueryError{QueryString: l.findEventsRowByRVQuery, Err: err}
+			}
+			defer rows.Close()
+
+			if !rows.Next() {
+				// query returned no results
+				if targetRV != latestRV {
+					return ErrTooOld
+				}
+				return nil
+			}
+			if err := rows.Scan(&rowID); err != nil {
+				return fmt.Errorf("failed scan rowid: %w", err)
 			}
-		} else if err != nil {
-			return fmt.Errorf("failed scan rowid: %w", err)
+			return nil
+		}(); err != nil {
+			return err
 		}
 
 		// Backfilling previous events from resourceVersion
-		rows, err := tx.Stmt(l.listEventsAfterStmt).QueryContext(ctx, rowID)
+		rows, err := l.QueryForRows(ctx, tx.Stmt(l.listEventsAfterStmt), rowID)
 		if err != nil {
 			return &db.QueryError{QueryString: l.listEventsAfterQuery, Err: err}
 		}
@@ -1027,9 +1036,8 @@ func (l *ListOptionIndexer) executeQuery(ctx context.Context, queryInfo *QueryIn
 
 	var items []any
 	err = l.WithTransaction(ctx, false, func(tx transaction.Client) error {
-		txStmt := tx.Stmt(stmt)
 		now := time.Now()
-		rows, err := txStmt.QueryContext(ctx, queryInfo.params...)
+		rows, err := l.QueryForRows(ctx, tx.Stmt(stmt), queryInfo.params...)
 		if err != nil {
 			return &db.QueryError{QueryString: queryInfo.query, Err: err}
 		}
@@ -1049,9 +1057,8 @@ func (l *ListOptionIndexer) executeQuery(ctx context.Context, queryInfo *QueryIn
 					err = errors.Join(err, &db.QueryError{QueryString: queryInfo.countQuery, Err: cerr})
 				}
 			}()
-			txStmt := tx.Stmt(countStmt)
 			now = time.Now()
-			rows, err := txStmt.QueryContext(ctx, queryInfo.countParams...)
+			rows, err := l.QueryForRows(ctx, tx.Stmt(countStmt), queryInfo.countParams...)
 			if err != nil {
 				return &db.QueryError{QueryString: queryInfo.countQuery, Err: err}
 			}