Skip to content

Commit 640dcf4

Browse files
koukou
authored
Fix out of range access with 1M over records response (#155)
Closes GH-154
1 parent 0176b4c commit 640dcf4

File tree

3 files changed

+75
-23
lines changed

3 files changed

+75
-23
lines changed

‎src/afs.cc

Lines changed: 26 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,40 +1032,42 @@ class ArrowArrayBuilder<ArrowType, arrow::enable_if_has_c_type<ArrowType>>
10321032

10331033
arrow::Status build_not_null(uint64_t iTuple, uint64_t iTupleEnd)
10341034
{
1035-
values_.resize(iTupleEnd - iTuple);
1036-
for (; iTuple < iTupleEnd; iTuple++)
1035+
uint64_t n = iTupleEnd - iTuple;
1036+
values_.resize(n);
1037+
for (uint64_t i = 0; i < n; ++i)
10371038
{
10381039
bool isNull;
1039-
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
1040+
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
10401041
SPI_tuptable->tupdesc,
10411042
iAttribute_ + 1,
10421043
&isNull);
1043-
values_[iTuple] = convert_value<ArrowType>(datum);
1044+
values_[i] = convert_value<ArrowType>(datum);
10441045
}
10451046
return builder_->AppendValues(values_.data(), values_.size());
10461047
}
10471048

10481049
arrow::Status build_may_null(uint64_t iTuple, uint64_t iTupleEnd)
10491050
{
10501051
bool haveNull = false;
1051-
values_.resize(iTupleEnd - iTuple);
1052-
validBytes_.resize(iTupleEnd - iTuple);
1053-
for (; iTuple < iTupleEnd; iTuple++)
1052+
uint64_t n = iTupleEnd - iTuple;
1053+
values_.resize(n);
1054+
validBytes_.resize(n);
1055+
for (uint64_t i = 0; i < n; ++i)
10541056
{
10551057
bool isNull;
1056-
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
1058+
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
10571059
SPI_tuptable->tupdesc,
10581060
iAttribute_ + 1,
10591061
&isNull);
10601062
if (isNull)
10611063
{
10621064
haveNull = true;
1063-
validBytes_[iTuple] = 0;
1065+
validBytes_[i] = 0;
10641066
}
10651067
else
10661068
{
1067-
validBytes_[iTuple] = 1;
1068-
values_[iTuple] = convert_value<ArrowType>(datum);
1069+
validBytes_[i] = 1;
1070+
values_[i] = convert_value<ArrowType>(datum);
10691071
}
10701072
}
10711073
if (haveNull)
@@ -1148,41 +1150,42 @@ class ArrowArrayBuilder<ArrowType, arrow::enable_if_base_binary<ArrowType>>
11481150

11491151
arrow::Status build_not_null(uint64_t iTuple, uint64_t iTupleEnd)
11501152
{
1151-
values_.resize(iTupleEnd - iTuple);
1152-
for (; iTuple < iTupleEnd; iTuple++)
1153+
uint64_t n = iTupleEnd - iTuple;
1154+
values_.resize(n);
1155+
for (uint64_t i = 0; i < n; ++i)
11531156
{
11541157
bool isNull;
1155-
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
1158+
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
11561159
SPI_tuptable->tupdesc,
11571160
iAttribute_ + 1,
11581161
&isNull);
1159-
values_[iTuple] = std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
1162+
values_[i] = std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
11601163
}
11611164
return builder_->AppendValues(values_);
11621165
}
11631166

11641167
arrow::Status build_may_null(uint64_t iTuple, uint64_t iTupleEnd)
11651168
{
11661169
bool haveNull = false;
1167-
values_.resize(iTupleEnd - iTuple);
1168-
validBytes_.resize(iTupleEnd - iTuple);
1169-
for (; iTuple < iTupleEnd; iTuple++)
1170+
uint64_t n = iTupleEnd - iTuple;
1171+
values_.resize(n);
1172+
validBytes_.resize(n);
1173+
for (uint64_t i = 0; i < n; ++i)
11701174
{
11711175
bool isNull;
1172-
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
1176+
auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
11731177
SPI_tuptable->tupdesc,
11741178
iAttribute_ + 1,
11751179
&isNull);
11761180
if (isNull)
11771181
{
11781182
haveNull = true;
1179-
validBytes_[iTuple] = 0;
1183+
validBytes_[i] = 0;
11801184
}
11811185
else
11821186
{
1183-
validBytes_[iTuple] = 1;
1184-
values_[iTuple] =
1185-
std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
1187+
validBytes_[i] = 1;
1188+
values_[i] = std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
11861189
}
11871190
}
11881191
if (haveNull)

‎test/helper/sandbox.rb

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ def running?
123123
end
124124

125125
def initdb(shared_preload_libraries: [],
126+
max_n_rows_per_record_batch: nil,
126127
db_path: "db",
127128
port: 25432,
128129
flight_sql_port: 35432)
@@ -163,6 +164,10 @@ def initdb(shared_preload_libraries: [],
163164
conf.puts("shared_preload_libraries = " +
164165
"'#{shared_preload_libraries.join(",")}'")
165166
conf.puts("arrow_flight_sql.uri = '#{@flight_sql_uri}'")
167+
if max_n_rows_per_record_batch
168+
conf.puts("arrow_flight_sql.max_n_rows_per_record_batch = " +
169+
"#{max_n_rows_per_record_batch}")
170+
end
166171
yield(conf) if block_given?
167172
end
168173
pg_hba_conf = File.join(@dir, "pg_hba.conf")
@@ -343,6 +348,7 @@ def setup_db
343348
@postgresql = PostgreSQL.new(@tmp_dir)
344349
options = {
345350
shared_preload_libraries: shared_preload_libraries,
351+
max_n_rows_per_record_batch: max_n_rows_per_record_batch,
346352
}
347353
@postgresql.initdb(**options)
348354
yield
@@ -352,6 +358,10 @@ def shared_preload_libraries
352358
["arrow_flight_sql"]
353359
end
354360

361+
def max_n_rows_per_record_batch
362+
nil
363+
end
364+
355365
def start_postgres
356366
@postgresql.start
357367
end

‎test/test-flight-sql.rb

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@
1818
class FlightSQLTest < Test::Unit::TestCase
1919
include Helper::Sandbox
2020

21+
def max_n_rows_per_record_batch
22+
10
23+
end
24+
2125
setup do
2226
unless flight_client.respond_to?(:authenticate_basic)
2327
omit("ArrowFlight::Client#authenticate_basic is needed" )
@@ -123,6 +127,41 @@ def test_select_from
123127
reader.read_all)
124128
end
125129

130+
data(:type, ["integer", "text"])
131+
data(:may_null, [true, false])
132+
def test_select_multiple_record_batches_response
133+
type = data[:type]
134+
may_null = data[:may_null]
135+
if type == "integer"
136+
arrow_type = Arrow::Int32DataType.new
137+
values = 30.times.to_a
138+
else
139+
arrow_type = Arrow::StringDataType.new
140+
values = 30.times.collect(&:to_s)
141+
end
142+
143+
create_table = "CREATE TABLE data (value #{type}"
144+
# XXX: "NOT NULL" doesn't change FormData_pg_attribute::attnotnull
145+
# returned by "SELECT *". So the "may_null" data driven test
146+
# parameter is meaningless...
147+
create_table << " NOT NULL" unless may_null
148+
create_table << ")"
149+
run_sql(create_table)
150+
run_sql("INSERT INTO data VALUES " +
151+
values.collect {|value| "(#{value})"}.join(", "))
152+
153+
info = flight_sql_client.execute("SELECT * FROM data", @options)
154+
assert_equal(Arrow::Schema.new(value: arrow_type),
155+
info.get_schema)
156+
endpoint = info.endpoints.first
157+
reader = flight_sql_client.do_get(endpoint.ticket, @options)
158+
chunks = values.each_slice(max_n_rows_per_record_batch).collect do |chunk|
159+
arrow_type.array_class.new(chunk)
160+
end
161+
assert_equal(Arrow::Table.new(value: Arrow::ChunkedArray.new(chunks)),
162+
reader.read_all)
163+
end
164+
126165
def test_insert_direct
127166
unless flight_sql_client.respond_to?(:execute_update)
128167
omit("red-arrow-flight-sql 13.0.0 or later is required")

0 commit comments

Comments
 (0)