Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

sql: fix iteration conditions in crdb_internal.scan #98329

Merged
merged 1 commit into from
Mar 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions pkg/sql/logictest/testdata/logic_test/sql_keys
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,23 @@ SELECT crdb_internal.pretty_key(key, 0) FROM crdb_internal.scan(crdb_internal.ta
/10/Table/106/1/1/0
/10/Table/106/1/2/0

# Regression test for a panic when the end key was equal to the Next() key of the last key returned.
statement ok
SELECT
crdb_internal.pretty_key(key, 0)
FROM
crdb_internal.scan(
ARRAY[
crdb_internal.table_span($tableid)[1],
(
SELECT key || '\x00'::BYTES
FROM crdb_internal.scan(crdb_internal.table_span($tableid))
ORDER BY key DESC
LIMIT 1
)
]
)

# An error should be returned when an invalid range ID is specified.
statement error pq: range with ID 1000000 not found
SELECT key FROM crdb_internal.list_sql_keys_in_range(1000000)
Expand Down
12 changes: 8 additions & 4 deletions pkg/sql/sem/builtins/generator_builtins.go
Original file line number Diff line number Diff line change
Expand Up @@ -2158,6 +2158,10 @@ type spanKeyIterator struct {
// the iterator maintains a small set of K/V pairs in the span,
// and accesses more in a streaming fashion.
kvs []roachpb.KeyValue

// resumeSpan is the resume span from the last ScanRequest.
resumeSpan *roachpb.Span

// index maintains the current position of the iterator in kvs.
index int
// A buffer to avoid allocating an array on every call to Values().
Expand Down Expand Up @@ -2189,15 +2193,14 @@ func (sp *spanKeyIterator) Next(ctx context.Context) (bool, error) {
return true, nil
}

// If we don't have any K/V pairs at all, then we're out of results.
if len(sp.kvs) == 0 {
// If we don't have a resume span, then we're out of results.
if sp.resumeSpan == nil {
return false, nil
}

// If we had some K/V pairs already, use the last key to constrain
// the result of the next scan.
startKey := sp.kvs[len(sp.kvs)-1].Key.Next()
err := sp.scan(ctx, startKey, sp.span.EndKey)
err := sp.scan(ctx, sp.resumeSpan.Key, sp.span.EndKey)
if err != nil {
return false, err
}
Expand All @@ -2224,6 +2227,7 @@ func (sp *spanKeyIterator) scan(
}
resp := br.Responses[0].GetScan()
sp.kvs = resp.Rows
sp.resumeSpan = resp.ResumeSpan
// The user of the generator first calls Next(), then Values(), so the index
// managing the iterator's position needs to start at -1 instead of 0.
sp.index = -1
Expand Down