lib/sql/src/statement/result.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const sql = @import("../root.zig");
  3 const access_mod = @import("access.zig");
  4 const ast_mod = @import("ast.zig");
  5 const execute_mod = @import("execute.zig");
  6 const catalog_mod = sql.catalog;
  7 const row = sql.row;
  8 const session_mod = sql.session;
  9 
 10 const Allocator = std.mem.Allocator;
 11 
 12 pub const Binding = struct {
 13     state: State = .unset,
 14     bytes: std.ArrayList(u8) = .empty,
 15 
 16     const State = union(enum) {
 17         unset,
 18         nil,
 19         integer: i64,
 20         text,
 21         blob,
 22     };
 23 
 24     pub const unset: Binding = .{};
 25 
 26     pub fn deinit(self: *Binding, allocator: Allocator) void {
 27         self.bytes.deinit(allocator);
 28         self.* = undefined;
 29     }
 30 
 31     pub fn clear(self: *Binding) void {
 32         self.state = .unset;
 33         self.bytes.clearRetainingCapacity();
 34     }
 35 
 36     pub fn set(self: *Binding, allocator: Allocator, bound_value: row.Value) Allocator.Error!void {
 37         switch (bound_value) {
 38             .nil => {
 39                 self.state = .nil;
 40                 self.bytes.clearRetainingCapacity();
 41             },
 42             .integer => |integer| {
 43                 self.state = .{ .integer = integer };
 44                 self.bytes.clearRetainingCapacity();
 45             },
 46             .text => |text| try self.setBytes(allocator, text, .text),
 47             .blob => |blob| try self.setBytes(allocator, blob, .blob),
 48         }
 49     }
 50 
 51     pub fn value(self: *const Binding) row.Value {
 52         return switch (self.state) {
 53             .unset, .nil => row.Value.nil,
 54             .integer => |integer| .{ .integer = integer },
 55             .text => .{ .text = self.bytes.items },
 56             .blob => .{ .blob = self.bytes.items },
 57         };
 58     }
 59 
 60     fn setBytes(self: *Binding, allocator: Allocator, source: []const u8, state: State) Allocator.Error!void {
 61         const offset = aliasOffset(self.bytes.items, source);
 62         try self.bytes.ensureTotalCapacity(allocator, source.len);
 63         const active_source = if (offset) |start| self.bytes.items[start..][0..source.len] else source;
 64         self.bytes.items.len = source.len;
 65         std.mem.copyForwards(u8, self.bytes.items, active_source);
 66         self.state = state;
 67     }
 68 
 69     fn aliasOffset(storage: []const u8, source: []const u8) ?usize {
 70         if (storage.len == 0 or source.len == 0) return null;
 71         const storage_start = @intFromPtr(storage.ptr);
 72         const storage_end = storage_start + storage.len;
 73         const source_start = @intFromPtr(source.ptr);
 74         const source_end = source_start + source.len;
 75         if (source_start < storage_start or source_end > storage_end) return null;
 76         return source_start - storage_start;
 77     }
 78 };
 79 
 80 test "Binding retains warmed text and blob storage" {
 81     var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});
 82     var binding = Binding.unset;
 83     defer binding.deinit(failing.allocator());
 84 
 85     try binding.set(failing.allocator(), .{ .text = "a retained binding payload" });
 86     const storage_ptr = binding.bytes.items.ptr;
 87     const storage_capacity = binding.bytes.capacity;
 88 
 89     failing.fail_index = failing.alloc_index;
 90     failing.resize_fail_index = failing.resize_index;
 91     for (0..8) |_| {
 92         try binding.set(failing.allocator(), .{ .text = "short text" });
 93         try std.testing.expectEqualStrings("short text", binding.value().text);
 94         try binding.set(failing.allocator(), .{ .blob = "blob bytes" });
 95         try std.testing.expectEqualSlices(u8, "blob bytes", binding.value().blob);
 96         try binding.set(failing.allocator(), .{ .integer = 42 });
 97         try std.testing.expectEqual(@as(i64, 42), binding.value().integer);
 98         try binding.set(failing.allocator(), .nil);
 99         try std.testing.expectEqual(row.Value.nil, binding.value());
100         binding.clear();
101         try std.testing.expectEqual(row.Value.nil, binding.value());
102     }
103 
104     try std.testing.expect(!failing.has_induced_failure);
105     try std.testing.expectEqual(storage_ptr, binding.bytes.items.ptr);
106     try std.testing.expectEqual(storage_capacity, binding.bytes.capacity);
107 }
108 
109 test "Binding preserves its value when growth fails" {
110     var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});
111     var binding = Binding.unset;
112     defer binding.deinit(failing.allocator());
113 
114     try binding.set(failing.allocator(), .{ .text = "original" });
115     const storage_ptr = binding.bytes.items.ptr;
116     const storage_capacity = binding.bytes.capacity;
117     var large: [4096]u8 = undefined;
118     for (&large, 0..) |*byte, index| byte.* = @intCast('a' + index % 26);
119 
120     failing.fail_index = failing.alloc_index;
121     failing.resize_fail_index = failing.resize_index;
122     try std.testing.expectError(error.OutOfMemory, binding.set(failing.allocator(), .{ .blob = &large }));
123     try std.testing.expectEqualStrings("original", binding.value().text);
124     try std.testing.expectEqual(storage_ptr, binding.bytes.items.ptr);
125     try std.testing.expectEqual(storage_capacity, binding.bytes.capacity);
126 
127     failing.fail_index = std.math.maxInt(usize);
128     failing.resize_fail_index = std.math.maxInt(usize);
129     try binding.set(failing.allocator(), .{ .blob = &large });
130     try std.testing.expectEqualSlices(u8, &large, binding.value().blob);
131 }
132 
133 test "Binding accepts a slice of its current value" {
134     var binding = Binding.unset;
135     defer binding.deinit(std.testing.allocator);
136 
137     try binding.set(std.testing.allocator, .{ .text = "prefix-suffix" });
138     const suffix = binding.value().text["prefix-".len..];
139     try binding.set(std.testing.allocator, .{ .text = suffix });
140     try std.testing.expectEqualStrings("suffix", binding.value().text);
141 }
142 
143 pub const RowStorage = union(enum) {
144     single: ?[]u8,
145     many: [][]u8,
146 };
147 
148 pub const Rows = struct {
149     allocator: Allocator,
150     storage: RowStorage,
151     cursor: usize = 0,
152 
153     pub fn deinit(self: *Rows) void {
154         switch (self.storage) {
155             .single => |bytes| if (bytes) |owned| self.allocator.free(owned),
156             .many => |rows| {
157                 for (rows) |bytes| self.allocator.free(bytes);
158                 self.allocator.free(rows);
159             },
160         }
161         self.* = undefined;
162     }
163 
164     pub fn next(self: *Rows) ?[]const u8 {
165         switch (self.storage) {
166             .single => |bytes| {
167                 if (self.cursor != 0) return null;
168                 self.cursor = 1;
169                 return bytes;
170             },
171             .many => |rows| {
172                 if (self.cursor >= rows.len) return null;
173                 const bytes = rows[self.cursor];
174                 self.cursor += 1;
175                 return bytes;
176             },
177         }
178     }
179 
180     pub fn count(self: *const Rows) usize {
181         return switch (self.storage) {
182             .single => |bytes| if (bytes == null) 0 else 1,
183             .many => |rows| rows.len,
184         };
185     }
186 };
187 
188 pub const Result = union(enum) {
189     mutation: session_mod.RelationFlush,
190     staged: usize,
191     catalog: session_mod.CatalogFlush,
192     rows: Rows,
193 
194     pub fn deinit(self: *Result, allocator: Allocator) void {
195         _ = allocator;
196         switch (self.*) {
197             .mutation => {},
198             .staged => {},
199             .catalog => {},
200             .rows => |*rows| rows.deinit(),
201         }
202         self.* = undefined;
203     }
204 
205     pub fn nextRow(self: *Result) ?[]const u8 {
206         return switch (self.*) {
207             .mutation => null,
208             .staged => null,
209             .catalog => null,
210             .rows => |*rows| rows.next(),
211         };
212     }
213 
214     pub fn rowCount(self: *const Result) usize {
215         return switch (self.*) {
216             .mutation => 0,
217             .staged => 0,
218             .catalog => 0,
219             .rows => |*rows| rows.count(),
220         };
221     }
222 };
223 
224 pub fn appendSelectedRow(result_allocator: Allocator, selected: *std.ArrayList([]u8), fields: []const execute_mod.ProjectedColumn, rowid: i64, bytes: []const u8) ast_mod.Error!void {
225     const owned = try selectedRow(result_allocator, fields, rowid, bytes);
226     errdefer result_allocator.free(owned);
227     try selected.append(result_allocator, owned);
228 }
229 
230 pub fn appendSelectedIndexRow(
231     result_allocator: Allocator,
232     selected: *std.ArrayList([]u8),
233     fields: []const execute_mod.ProjectedColumn,
234     handle: *const catalog_mod.RelationHandle,
235     index_slot: usize,
236     index_values: []const row.Value,
237     rowid: i64,
238 ) ast_mod.Error!void {
239     const owned = try selectedIndexRow(result_allocator, fields, handle, index_slot, index_values, rowid);
240     errdefer result_allocator.free(owned);
241     try selected.append(result_allocator, owned);
242 }
243 
244 pub fn selectedRow(result_allocator: Allocator, fields: []const execute_mod.ProjectedColumn, rowid: i64, bytes: []const u8) ast_mod.Error![]u8 {
245     if (fields.len == 0) return try result_allocator.dupe(u8, bytes);
246     const view = try row.View.init(bytes);
247     var stack_values: [catalog_mod.max_columns + 1]row.Value = undefined;
248     const stack_backed = fields.len <= stack_values.len;
249     const values = if (stack_backed) stack_values[0..fields.len] else try result_allocator.alloc(row.Value, fields.len);
250     defer if (!stack_backed) result_allocator.free(values);
251     for (fields, 0..) |field, offset| {
252         values[offset] = switch (field) {
253             .rowid => .{ .integer = rowid },
254             .field => |index| try view.column(index),
255         };
256     }
257     const size = try row.encodedSize(values);
258     const projected = try result_allocator.alloc(u8, size);
259     errdefer result_allocator.free(projected);
260     _ = try row.encode(projected, values);
261     return projected;
262 }
263 
264 pub fn selectedRowInto(result_allocator: Allocator, row_buffer: *std.ArrayList(u8), fields: []const execute_mod.ProjectedColumn, rowid: i64, bytes: []const u8) ast_mod.Error![]const u8 {
265     if (fields.len == 0) {
266         try row_buffer.resize(result_allocator, bytes.len);
267         @memcpy(row_buffer.items, bytes);
268         return row_buffer.items;
269     }
270     const view = try row.View.init(bytes);
271     var stack_values: [catalog_mod.max_columns + 1]row.Value = undefined;
272     const stack_backed = fields.len <= stack_values.len;
273     const values = if (stack_backed) stack_values[0..fields.len] else try result_allocator.alloc(row.Value, fields.len);
274     defer if (!stack_backed) result_allocator.free(values);
275     for (fields, 0..) |field, offset| {
276         values[offset] = switch (field) {
277             .rowid => .{ .integer = rowid },
278             .field => |index| try view.column(index),
279         };
280     }
281     const size = try row.encodedSize(values);
282     try row_buffer.resize(result_allocator, size);
283     _ = try row.encode(row_buffer.items, values);
284     return row_buffer.items;
285 }
286 
287 pub fn selectedIndexRow(
288     result_allocator: Allocator,
289     fields: []const execute_mod.ProjectedColumn,
290     handle: *const catalog_mod.RelationHandle,
291     index_slot: usize,
292     index_values: []const row.Value,
293     rowid: i64,
294 ) ast_mod.Error![]u8 {
295     const spec = handle.specs[index_slot];
296     var stack_values: [catalog_mod.max_columns + 1]row.Value = undefined;
297     const stack_backed = fields.len <= stack_values.len;
298     const values = if (stack_backed) stack_values[0..fields.len] else try result_allocator.alloc(row.Value, fields.len);
299     defer if (!stack_backed) result_allocator.free(values);
300     for (fields, 0..) |field, offset| {
301         values[offset] = switch (field) {
302             .rowid => .{ .integer = rowid },
303             .field => |index| index_values[access_mod.fieldPosition(spec.fields, index) orelse unreachable],
304         };
305     }
306     const size = try row.encodedSize(values);
307     const projected = try result_allocator.alloc(u8, size);
308     errdefer result_allocator.free(projected);
309     _ = try row.encode(projected, values);
310     return projected;
311 }
312 
313 pub fn selectedIndexRowInto(
314     result_allocator: Allocator,
315     row_buffer: *std.ArrayList(u8),
316     fields: []const execute_mod.ProjectedColumn,
317     handle: *const catalog_mod.RelationHandle,
318     index_slot: usize,
319     index_values: []const row.Value,
320     rowid: i64,
321 ) ast_mod.Error![]const u8 {
322     const spec = handle.specs[index_slot];
323     var stack_values: [catalog_mod.max_columns + 1]row.Value = undefined;
324     const stack_backed = fields.len <= stack_values.len;
325     const values = if (stack_backed) stack_values[0..fields.len] else try result_allocator.alloc(row.Value, fields.len);
326     defer if (!stack_backed) result_allocator.free(values);
327     for (fields, 0..) |field, offset| {
328         values[offset] = switch (field) {
329             .rowid => .{ .integer = rowid },
330             .field => |index| index_values[access_mod.fieldPosition(spec.fields, index) orelse unreachable],
331         };
332     }
333     const size = try row.encodedSize(values);
334     try row_buffer.resize(result_allocator, size);
335     _ = try row.encode(row_buffer.items, values);
336     return row_buffer.items;
337 }
338 
339 pub fn freeSelectedRows(result_allocator: Allocator, rows: []const []u8) void {
340     for (rows) |bytes| result_allocator.free(bytes);
341 }
342 
343 pub const SortedRow = struct {
344     rowid: i64,
345     projected: []u8,
346     keys: []const row.Value,
347 };
348 
349 pub const SortedRowContext = struct {
350     order_fields: []const execute_mod.OrderField,
351 };
352 
353 pub fn orderKeyCount(order_fields: []const execute_mod.OrderField) usize {
354     var count: usize = 0;
355     for (order_fields) |order_field| {
356         if (order_field.source == .field) count += 1;
357     }
358     return count;
359 }
360 
361 pub fn appendSortedRow(
362     result_allocator: Allocator,
363     collected: *std.ArrayList(SortedRow),
364     fields: []const execute_mod.ProjectedColumn,
365     order_fields: []const execute_mod.OrderField,
366     rowid: i64,
367     bytes: []const u8,
368 ) ast_mod.Error!void {
369     const view = try row.View.init(bytes);
370     const key_count = orderKeyCount(order_fields);
371     const keys: []row.Value = if (key_count == 0)
372         &.{}
373     else
374         try result_allocator.alloc(row.Value, key_count);
375     var extracted: usize = 0;
376     errdefer {
377         for (keys[0..extracted]) |key_value| freeValue(result_allocator, key_value);
378         result_allocator.free(keys);
379     }
380     for (order_fields) |order_field| {
381         switch (order_field.source) {
382             .rowid => {},
383             .field => |field| {
384                 keys[extracted] = try copyValue(result_allocator, sortedColumn(view, field));
385                 extracted += 1;
386             },
387         }
388     }
389     std.debug.assert(extracted == key_count);
390     const projected = try selectedRow(result_allocator, fields, rowid, bytes);
391     errdefer result_allocator.free(projected);
392     try collected.append(result_allocator, .{
393         .rowid = rowid,
394         .projected = projected,
395         .keys = keys,
396     });
397 }
398 
399 pub fn appendSortedIndexRow(
400     result_allocator: Allocator,
401     collected: *std.ArrayList(SortedRow),
402     fields: []const execute_mod.ProjectedColumn,
403     order_fields: []const execute_mod.OrderField,
404     handle: *const catalog_mod.RelationHandle,
405     index_slot: usize,
406     index_values: []const row.Value,
407     rowid: i64,
408 ) ast_mod.Error!void {
409     const spec = handle.specs[index_slot];
410     const key_count = orderKeyCount(order_fields);
411     const keys: []row.Value = if (key_count == 0)
412         &.{}
413     else
414         try result_allocator.alloc(row.Value, key_count);
415     var extracted: usize = 0;
416     errdefer {
417         for (keys[0..extracted]) |key_value| freeValue(result_allocator, key_value);
418         result_allocator.free(keys);
419     }
420     for (order_fields) |order_field| {
421         switch (order_field.source) {
422             .rowid => {},
423             .field => |field| {
424                 const position = access_mod.fieldPosition(spec.fields, field) orelse unreachable;
425                 keys[extracted] = try copyValue(result_allocator, index_values[position]);
426                 extracted += 1;
427             },
428         }
429     }
430     std.debug.assert(extracted == key_count);
431     const projected = try selectedIndexRow(
432         result_allocator,
433         fields,
434         handle,
435         index_slot,
436         index_values,
437         rowid,
438     );
439     errdefer result_allocator.free(projected);
440     try collected.append(result_allocator, .{
441         .rowid = rowid,
442         .projected = projected,
443         .keys = keys,
444     });
445 }
446 
447 pub fn freeSortedRowContents(result_allocator: Allocator, rows: []const SortedRow) void {
448     for (rows) |sorted_row| {
449         result_allocator.free(sorted_row.projected);
450         freeSortedKeys(result_allocator, sorted_row.keys);
451     }
452 }
453 
454 pub fn freeSortedKeys(result_allocator: Allocator, keys: []const row.Value) void {
455     for (keys) |key_value| freeValue(result_allocator, key_value);
456     result_allocator.free(keys);
457 }
458 
459 pub fn freeSortedRows(result_allocator: Allocator, rows: []SortedRow) void {
460     freeSortedRowContents(result_allocator, rows);
461     result_allocator.free(rows);
462 }
463 
464 pub fn sortedRowLessThan(context: SortedRowContext, left: SortedRow, right: SortedRow) bool {
465     return sortedRowOrder(context, left, right) == .lt;
466 }
467 
468 pub fn sortedRowOrder(context: SortedRowContext, left: SortedRow, right: SortedRow) std.math.Order {
469     std.debug.assert(left.keys.len == right.keys.len);
470     var key_index: usize = 0;
471     for (context.order_fields) |order_field| {
472         const order = switch (order_field.source) {
473             .rowid => std.math.order(left.rowid, right.rowid),
474             .field => blk: {
475                 std.debug.assert(key_index < left.keys.len);
476                 const compared = row.compareValues(
477                     left.keys[key_index],
478                     right.keys[key_index],
479                     order_field.collation,
480                 );
481                 key_index += 1;
482                 break :blk compared;
483             },
484         };
485         if (order != .eq) return if (order_field.descending) order.invert() else order;
486     }
487     return std.math.order(left.rowid, right.rowid);
488 }
489 
490 pub fn sortedColumn(view: row.View, field: usize) row.Value {
491     return view.column(field) catch |err| switch (err) {
492         error.ColumnOutOfBounds => .nil,
493         else => unreachable,
494     };
495 }
496 
497 pub fn copyValue(allocator: Allocator, value: row.Value) std.mem.Allocator.Error!row.Value {
498     return switch (value) {
499         .nil => row.Value.nil,
500         .integer => |integer| .{ .integer = integer },
501         .text => |text| .{ .text = try allocator.dupe(u8, text) },
502         .blob => |blob| .{ .blob = try allocator.dupe(u8, blob) },
503     };
504 }
505 
506 pub fn freeValue(allocator: Allocator, value: row.Value) void {
507     switch (value) {
508         .nil, .integer => {},
509         .text => |text| allocator.free(text),
510         .blob => |blob| allocator.free(blob),
511     }
512 }