From 217a0823c49474f56bbaa477fbd479c9b1c32d1a Mon Sep 17 00:00:00 2001 From: Vincent Rischmann Date: Thu, 12 Nov 2020 16:01:40 +0100 Subject: add an iterator An iterator can be used to process rows one by one, without loading everything into an array list first. --- sqlite.zig | 359 ++++++++++++++++++++++++++++++++++++------------------------- 1 file changed, 211 insertions(+), 148 deletions(-) (limited to 'sqlite.zig') diff --git a/sqlite.zig b/sqlite.zig index d08717b..c37b972 100644 --- a/sqlite.zig +++ b/sqlite.zig @@ -120,6 +120,139 @@ pub const Db = struct { } }; +pub fn Iterator(comptime Type: type) type { + return struct { + const Self = @This(); + + const TypeInfo = @typeInfo(Type); + + stmt: *c.sqlite3_stmt, + + pub fn next(self: *Self, options: anytype) !?Type { + var result = c.sqlite3_step(self.stmt); + if (result == c.SQLITE_DONE) { + return null; + } + + if (result != c.SQLITE_ROW) { + logger.err("unable to iterate, result: {}", .{result}); + return error.SQLiteStepError; + } + + const columns = c.sqlite3_column_count(self.stmt); + + return switch (TypeInfo) { + .Int => blk: { + debug.assert(columns == 1); + break :blk try self.readInt(options); + }, + .Float => blk: { + debug.assert(columns == 1); + break :blk try self.readFloat(options); + }, + .Struct => blk: { + std.debug.assert(columns == TypeInfo.Struct.fields.len); + break :blk try self.readStruct(options); + }, + else => @compileError("cannot read into type " ++ @typeName(Type)), + }; + } + + fn readInt(self: *Self, options: anytype) !Type { + const n = c.sqlite3_column_int64(self.stmt, 0); + return @intCast(Type, n); + } + + fn readFloat(self: *Self, options: anytype) !Type { + const d = c.sqlite3_column_double(self.stmt, 0); + return @floatCast(Type, d); + } + + const ReadBytesMode = enum { + Blob, + Text, + }; + + fn readBytes(self: *Self, options: anytype, mode: ReadBytesMode, _i: usize, ptr: *[]const u8) !void { + const i = @intCast(c_int, _i); + switch (mode) { + .Blob => { + const data = c.sqlite3_column_blob(self.stmt, i); + if (data == null) ptr.* = ""; + + const size = @intCast(usize, c.sqlite3_column_bytes(self.stmt, i)); + + var tmp = try options.allocator.alloc(u8, size); + mem.copy(u8, tmp, @ptrCast([*c]const u8, data)[0..size]); + + ptr.* = tmp; + }, + .Text => { + const data = c.sqlite3_column_text(self.stmt, i); + if (data == null) ptr.* = ""; + + const size = @intCast(usize, c.sqlite3_column_bytes(self.stmt, i)); + + var tmp = try options.allocator.alloc(u8, size); + mem.copy(u8, tmp, @ptrCast([*c]const u8, data)[0..size]); + + ptr.* = tmp; + }, + } + } + + fn readStruct(self: *Self, options: anytype) !Type { + var value: Type = undefined; + + inline for (@typeInfo(Type).Struct.fields) |field, _i| { + const i = @as(usize, _i); + const field_type_info = @typeInfo(field.field_type); + + switch (field.field_type) { + []const u8, []u8 => { + try self.readBytes(options, .Blob, i, &@field(value, field.name)); + }, + Blob => { + try self.readBytes(options, .Blob, i, &@field(value, field.name).data); + }, + Text => { + try self.readBytes(options, .Text, i, &@field(value, field.name).data); + }, + else => switch (field_type_info) { + .Int => { + const n = c.sqlite3_column_int64(self.stmt, i); + @field(value, field.name) = @intCast(field.field_type, n); + }, + .Float => { + const f = c.sqlite3_column_double(self.stmt, i); + @field(value, field.name) = f; + }, + .Void => { + @field(value, field.name) = {}; + }, + .Array => |arr| { + switch (arr.child) { + u8 => { + const data = c.sqlite3_column_blob(self.stmt, i); + const size = @intCast(usize, c.sqlite3_column_bytes(self.stmt, i)); + + if (size > @as(usize, arr.len)) return error.ArrayTooSmall; + + mem.copy(u8, @field(value, field.name)[0..], @ptrCast([*c]const u8, data)[0..size]); + }, + else => @compileError("cannot populate field " ++ field.name ++ " of type array of " ++ @typeName(arr.child)), + } + }, + else => @compileError("cannot populate field " ++ field.name ++ " of type " ++ @typeName(field.field_type)), + }, + } + } + + return value; + } + }; +} + pub const StatementOptions = struct {}; /// Statement is a wrapper around a SQLite statement, providing high-level functions to execute @@ -222,7 +355,7 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t /// } /// /// The types are checked at comptime. - pub fn bind(self: *Self, values: anytype) void { + fn bind(self: *Self, values: anytype) void { const StructType = @TypeOf(values); const StructTypeInfo = @typeInfo(StructType).Struct; @@ -285,6 +418,15 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t } } + pub fn iterator(self: *Self, comptime Type: type, values: anytype) !Iterator(Type) { + self.bind(values); + + var res: Iterator(Type) = undefined; + res.stmt = self.stmt; + + return res; + } + /// one reads a single row from the result set of this statement. /// /// The data in the row is used to populate a value of the type `Type`. @@ -312,32 +454,13 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t /// pub fn one(self: *Self, comptime Type: type, options: anytype, values: anytype) !?Type { if (!comptime std.meta.trait.is(.Struct)(@TypeOf(options))) { - @compileError("options passed to all must be a struct"); + @compileError("options passed to iterator must be a struct"); } - const TypeInfo = @typeInfo(Type); - - self.bind(values); - var result = c.sqlite3_step(self.stmt); + var iter = try self.iterator(Type, values); - switch (TypeInfo) { - .Int => return switch (result) { - c.SQLITE_ROW => try self.readInt(Type, options), - c.SQLITE_DONE => null, - else => std.debug.panic("invalid result {}", .{result}), - }, - .Float => return switch (result) { - c.SQLITE_ROW => try self.readFloat(Type, options), - c.SQLITE_DONE => null, - else => std.debug.panic("invalid result {}", .{result}), - }, - .Struct => return switch (result) { - c.SQLITE_ROW => try self.readStruct(Type, options), - c.SQLITE_DONE => null, - else => std.debug.panic("invalid result {}", .{result}), - }, - else => @compileError("cannot read into type " ++ @typeName(Type)), - } + const row = (try iter.next(options)) orelse return null; + return row; } /// all reads all rows from the result set of this statement. @@ -370,137 +493,18 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t /// pub fn all(self: *Self, comptime Type: type, options: anytype, values: anytype) ![]Type { if (!comptime std.meta.trait.is(.Struct)(@TypeOf(options))) { - @compileError("options passed to all must be a struct"); + @compileError("options passed to iterator must be a struct"); } - const TypeInfo = @typeInfo(Type); - - self.bind(values); + var iter = try self.iterator(Type, values); var rows = std.ArrayList(Type).init(options.allocator); - - var result = c.sqlite3_step(self.stmt); - while (result == c.SQLITE_ROW) : (result = c.sqlite3_step(self.stmt)) { - const columns = c.sqlite3_column_count(self.stmt); - - var value = switch (TypeInfo) { - .Int => blk: { - debug.assert(columns == 1); - break :blk try self.readInt(Type, options); - }, - .Float => blk: { - debug.assert(columns == 1); - break :blk try self.readFloat(Type, options); - }, - .Struct => blk: { - std.debug.assert(columns == @typeInfo(Type).Struct.fields.len); - break :blk try self.readStruct(Type, options); - }, - else => @compileError("cannot read into type " ++ @typeName(Type)), - }; - - try rows.append(value); - } - - if (result != c.SQLITE_DONE) { - logger.err("unable to iterate, result: {}", .{result}); - return error.SQLiteStepError; + while (true) { + const row = (try iter.next(options)) orelse break; + try rows.append(row); } return rows.span(); } - - fn readInt(self: *Self, comptime Type: type, options: anytype) !Type { - const n = c.sqlite3_column_int64(self.stmt, 0); - return @intCast(Type, n); - } - - fn readFloat(self: *Self, comptime Type: type, options: anytype) !Type { - const d = c.sqlite3_column_double(self.stmt, 0); - return @floatCast(Type, d); - } - - const ReadBytesMode = enum { - Blob, - Text, - }; - - fn readBytes(self: *Self, allocator: *mem.Allocator, mode: ReadBytesMode, _i: usize, ptr: *[]const u8) !void { - const i = @intCast(c_int, _i); - switch (mode) { - .Blob => { - const data = c.sqlite3_column_blob(self.stmt, i); - if (data == null) ptr.* = ""; - - const size = @intCast(usize, c.sqlite3_column_bytes(self.stmt, i)); - - var tmp = try allocator.alloc(u8, size); - mem.copy(u8, tmp, @ptrCast([*c]const u8, data)[0..size]); - - ptr.* = tmp; - }, - .Text => { - const data = c.sqlite3_column_text(self.stmt, i); - if (data == null) ptr.* = ""; - - const size = @intCast(usize, c.sqlite3_column_bytes(self.stmt, i)); - - var tmp = try allocator.alloc(u8, size); - mem.copy(u8, tmp, @ptrCast([*c]const u8, data)[0..size]); - - ptr.* = tmp; - }, - } - } - - fn readStruct(self: *Self, comptime Type: type, options: anytype) !Type { - var value: Type = undefined; - - inline for (@typeInfo(Type).Struct.fields) |field, _i| { - const i = @as(usize, _i); - const field_type_info = @typeInfo(field.field_type); - - switch (field.field_type) { - []const u8, []u8 => { - try self.readBytes(options.allocator, .Blob, i, &@field(value, field.name)); - }, - Blob => { - try self.readBytes(options.allocator, .Blob, i, &@field(value, field.name).data); - }, - Text => { - try self.readBytes(options.allocator, .Text, i, &@field(value, field.name).data); - }, - else => switch (field_type_info) { - .Int => { - const n = c.sqlite3_column_int64(self.stmt, i); - @field(value, field.name) = @intCast(field.field_type, n); - }, - .Float => { - const f = c.sqlite3_column_double(self.stmt, i); - @field(value, field.name) = f; - }, - .Void => { - @field(value, field.name) = {}; - }, - .Array => |arr| { - switch (arr.child) { - u8 => { - const data = c.sqlite3_column_blob(self.stmt, i); - const size = @intCast(usize, c.sqlite3_column_bytes(self.stmt, i)); - - if (size > @as(usize, arr.len)) return error.ArrayTooSmall; - - mem.copy(u8, @field(value, field.name)[0..], @ptrCast([*c]const u8, data)[0..size]); - }, - else => @compileError("cannot populate field " ++ field.name ++ " of type array of " ++ @typeName(arr.child)), - } - }, - else => @compileError("cannot populate field " ++ field.name ++ " of type " ++ @typeName(field.field_type)), - }, - } - } - - return value; - } }; } @@ -715,6 +719,65 @@ test "sqlite: statement reset" { } } +test "sqlite: statement iterator" { + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + var allocator = &arena.allocator; + + var db: Db = undefined; + try db.init(testing.allocator, .{ .mode = dbMode() }); + + // Create the tables + inline for (AllDDL) |ddl| { + try db.exec(ddl, .{}); + } + + // Add data + var stmt = try db.prepare("INSERT INTO user(id, name, age) VALUES(?{usize}, ?{[]const u8}, ?{usize})"); + defer stmt.deinit(); + + var expected_rows = std.ArrayList(TestUser).init(allocator); + var i: usize = 0; + while (i < 20) : (i += 1) { + const name = try std.fmt.allocPrint(allocator, "Vincent {}", .{i}); + const user = TestUser{ .id = i, .name = name, .age = i + 200 }; + + try expected_rows.append(user); + + stmt.reset(); + try stmt.exec(user); + + const rows_inserted = db.rowsAffected(); + testing.expectEqual(@as(usize, 1), rows_inserted); + } + + // Get the data with an iterator + var stmt2 = try db.prepare("SELECT name, age FROM user"); + defer stmt2.deinit(); + + const Type = struct { + name: Text, + age: usize, + }; + + var iter = try stmt2.iterator(Type, .{}); + + var rows = std.ArrayList(Type).init(allocator); + while (true) { + const row = (try iter.next(.{ .allocator = allocator })) orelse break; + try rows.append(row); + } + + // Check the data + testing.expectEqual(expected_rows.span().len, rows.span().len); + + for (rows.span()) |row, j| { + const exp_row = expected_rows.span()[j]; + testing.expectEqualStrings(exp_row.name, row.name.data); + testing.expectEqual(exp_row.age, row.age); + } +} + fn dbMode() Db.Mode { return if (build_options.is_ci) blk: { break :blk .{ .Memory = {} }; -- cgit v1.2.3 From 08919c71c87790a92a209e1c49bd76cbbcb2a1e3 Mon Sep 17 00:00:00 2001 From: Vincent Rischmann Date: Thu, 12 Nov 2020 16:44:23 +0100 Subject: add documentation for the iterator --- sqlite.zig | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) (limited to 'sqlite.zig') diff --git a/sqlite.zig b/sqlite.zig index c37b972..3c2e924 100644 --- a/sqlite.zig +++ b/sqlite.zig @@ -120,6 +120,28 @@ pub const Db = struct { } }; +/// Iterator allows iterating over a result set. +/// +/// Each call to `next` returns the next row of the result set, or null if the result set is exhausted. +/// Each row will have the type `Type` so the columns returned in the result set must be compatible with this type. +/// +/// Here is an example of how to use the iterator: +/// +/// const User = struct { +/// name: Text, +/// age: u16, +/// }; +/// +/// var stmt = try db.prepare("SELECT name, age FROM user"); +/// defer stmt.deinit(); +/// +/// var iter = try stmt.iterator(User, .{}); +/// while (true) { +/// const row: User = (try iter.next(.{})) orelse break; +/// ... +/// } +/// +/// The iterator _must not_ outlive the statement. pub fn Iterator(comptime Type: type) type { return struct { const Self = @This(); @@ -418,6 +440,25 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t } } + /// iterator returns an iterator to read data from the result set, one row at a time. + /// + /// The data in the row is used to populate a value of the type `Type`. + /// This means that `Type` must have as many fields as is returned in the query + /// executed by this statement. + /// This also means that the type of each field must be compatible with the SQLite type. + /// + /// Here is an example of how to use the iterator: + /// + /// var iter = try stmt.iterator(usize, .{}); + /// while (true) { + /// const row = (try iter.next(.{})) orelse break; + /// ... + /// } + /// + /// The `values` tuple is used for the bind parameters. It must have as many fields as there are bind markers + /// in the input query string. + /// + /// The iterator _must not_ outlive the statement. pub fn iterator(self: *Self, comptime Type: type, values: anytype) !Iterator(Type) { self.bind(values); -- cgit v1.2.3