1 files changed, 109 insertions, 48 deletions

diff --git a/sqlite.zig b/sqlite.zig
index e72c2d5..4cd91fd 100644
--- a/sqlite.zig
+++ b/sqlite.zig
@@ -203,6 +203,13 @@ pub const Db = struct {
         return try stmt.one(Type, options, values);
     }
 
+    /// oneAlloc is like `one` but can allocate memory.
+    pub fn oneAlloc(self: *Self, comptime Type: type, allocator: *mem.Allocator, comptime query: []const u8, options: anytype, values: anytype) !?Type {
+        var stmt = try self.prepare(query);
+        defer stmt.deinit();
+        return try stmt.oneAlloc(Type, allocator, options, values);
+    }
+
     /// prepare prepares a statement for the `query` provided.
     ///
     /// The query is analysed at comptime to search for bind markers.
@@ -259,16 +266,55 @@ pub fn Iterator(comptime Type: type) type {
         stmt: *c.sqlite3_stmt,
 
         // next scans the next row using the prepared statement.
-        //
         // If it returns null iterating is done.
+        //
+        // This cannot allocate memory. If you need to read TEXT or BLOB columns you need to use arrays or alternatively call nextAlloc.
         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) {
+                return error.SQLiteStepError;
+            }
 
+            const columns = c.sqlite3_column_count(self.stmt);
+
+            switch (TypeInfo) {
+                .Int => {
+                    debug.assert(columns == 1);
+                    return try self.readInt(Type, 0);
+                },
+                .Float => {
+                    debug.assert(columns == 1);
+                    return try self.readFloat(Type, 0);
+                },
+                .Bool => {
+                    debug.assert(columns == 1);
+                    return try self.readBool(0);
+                },
+                .Void => {
+                    debug.assert(columns == 1);
+                },
+                .Array => {
+                    debug.assert(columns == 1);
+                    return try self.readArray(Type, 0);
+                },
+                .Struct => {
+                    std.debug.assert(columns == TypeInfo.Struct.fields.len);
+                    return try self.readStruct(.{});
+                },
+                else => @compileError("cannot read into type " ++ @typeName(Type) ++ " ; if dynamic memory allocation is required use nextAlloc"),
+            }
+        }
+
+        // nextAlloc is like `next` but can allocate memory.
+        pub fn nextAlloc(self: *Self, allocator: *mem.Allocator, 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;
             }
 
@@ -277,15 +323,15 @@ pub fn Iterator(comptime Type: type) type {
             switch (Type) {
                 []const u8, []u8 => {
                     debug.assert(columns == 1);
-                    return try self.readBytes(Type, options.allocator, 0, .Text);
+                    return try self.readBytes(Type, allocator, 0, .Text);
                 },
                 Blob => {
                     debug.assert(columns == 1);
-                    return try self.readBytes(Blob, options.allocator, 0, .Blob);
+                    return try self.readBytes(Blob, allocator, 0, .Blob);
                 },
                 Text => {
                     debug.assert(columns == 1);
-                    return try self.readBytes(Text, options.allocator, 0, .Text);
+                    return try self.readBytes(Text, allocator, 0, .Text);
                 },
                 else => {},
             }
@@ -312,11 +358,13 @@ pub fn Iterator(comptime Type: type) type {
                 },
                 .Pointer => {
                     debug.assert(columns == 1);
-                    return try self.readPointer(Type, 0, options);
+                    return try self.readPointer(Type, allocator, 0);
                 },
                 .Struct => {
                     std.debug.assert(columns == TypeInfo.Struct.fields.len);
-                    return try self.readStruct(options);
+                    return try self.readStruct(.{
+                        .allocator = allocator,
+                    });
                 },
                 else => @compileError("cannot read into type " ++ @typeName(Type)),
             }
@@ -458,7 +506,7 @@ pub fn Iterator(comptime Type: type) type {
             }
         }
 
-        fn readPointer(self: *Self, comptime PointerType: type, i: usize, options: anytype) !PointerType {
+        fn readPointer(self: *Self, comptime PointerType: type, allocator: *mem.Allocator, i: usize) !PointerType {
             const type_info = @typeInfo(PointerType);
 
             var ret: PointerType = undefined;
@@ -467,7 +515,7 @@ pub fn Iterator(comptime Type: type) type {
                     switch (ptr.size) {
                         .One => unreachable,
                         .Slice => switch (ptr.child) {
-                            u8 => ret = try self.readBytes(PointerType, options.allocator, i, .Text),
+                            u8 => ret = try self.readBytes(PointerType, allocator, i, .Text),
                             else => @compileError("cannot read pointer of type " ++ @typeName(PointerType)),
                         },
                         else => @compileError("cannot read pointer of type " ++ @typeName(PointerType)),
@@ -516,7 +564,7 @@ pub fn Iterator(comptime Type: type) type {
                         .Bool => try self.readBool(i),
                         .Void => {},
                         .Array => try self.readArray(field.field_type, i),
-                        .Pointer => try self.readPointer(field.field_type, i, options),
+                        .Pointer => try self.readPointer(field.field_type, options.allocator, i),
                         else => @compileError("cannot populate field " ++ field.name ++ " of type " ++ @typeName(field.field_type)),
                     },
                 };
@@ -734,10 +782,10 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t
         ///     const row = try stmt.one(
         ///         struct {
         ///             id: usize,
-        ///             name: []const u8,
+        ///             name: [400]u8,
         ///             age: usize,
         ///         },
-        ///         .{ .allocator = allocator },
+        ///         .{},
         ///         .{ .foo = "bar", .age = 500 },
         ///     );
         ///
@@ -747,6 +795,7 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t
         /// 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.
         ///
+        /// This cannot allocate memory. If you need to read TEXT or BLOB columns you need to use arrays or alternatively call `oneAlloc`.
         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 iterator must be a struct");
@@ -758,6 +807,18 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t
             return row;
         }
 
+        /// oneAlloc is like `one` but can allocate memory.
+        pub fn oneAlloc(self: *Self, comptime Type: type, allocator: *mem.Allocator, options: anytype, values: anytype) !?Type {
+            if (!comptime std.meta.trait.is(.Struct)(@TypeOf(options))) {
+                @compileError("options passed to iterator must be a struct");
+            }
+
+            var iter = try self.iterator(Type, values);
+
+            const row = (try iter.nextAlloc(allocator, options)) orelse return null;
+            return row;
+        }
+
         /// all reads all rows from the result set of this statement.
         ///
         /// The data in each row is used to populate a value of the type `Type`.
@@ -773,7 +834,8 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t
         ///             name: []const u8,
         ///             age: usize,
         ///         },
-        ///         .{ .allocator = allocator },
+        ///         allocator,
+        ///         .{},
         ///         .{ .foo = "bar", .age = 500 },
         ///     );
         ///
@@ -783,18 +845,16 @@ pub fn Statement(comptime opts: StatementOptions, comptime query: ParsedQuery) t
         /// 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.
         ///
-        /// Note that this allocates all rows into a single slice: if you read a lot of data this can
-        /// use a lot of memory.
-        ///
-        pub fn all(self: *Self, comptime Type: type, options: anytype, values: anytype) ![]Type {
+        /// Note that this allocates all rows into a single slice: if you read a lot of data this can use a lot of memory.
+        pub fn all(self: *Self, comptime Type: type, allocator: *mem.Allocator, options: anytype, values: anytype) ![]Type {
             if (!comptime std.meta.trait.is(.Struct)(@TypeOf(options))) {
                 @compileError("options passed to iterator must be a struct");
             }
             var iter = try self.iterator(Type, values);
 
-            var rows = std.ArrayList(Type).init(options.allocator);
+            var rows = std.ArrayList(Type).init(allocator);
             while (true) {
-                const row = (try iter.next(options)) orelse break;
+                const row = (try iter.nextAlloc(allocator, options)) orelse break;
                 try rows.append(row);
             }
 
@@ -866,22 +926,25 @@ test "sqlite: db pragma" {
 
     if (build_options.in_memory) {
         const journal_mode = try db.pragma(
-            []const u8,
+            [128:0]u8,
             "journal_mode",
-            .{ .allocator = &arena.allocator },
+            .{},
             .{"wal"},
         );
         testing.expect(journal_mode != null);
         testing.expectEqualStrings("memory", journal_mode.?);
     } else {
-        const journal_mode = try db.pragma(
-            []const u8,
-            "journal_mode",
-            .{ .allocator = &arena.allocator },
-            .{"wal"},
-        );
-        testing.expect(journal_mode != null);
-        testing.expectEqualStrings("wal", journal_mode.?);
+        {
+            const journal_mode = try db.pragma([128:0]u8, .{}, "journal_mode", arg);
+            testing.expect(journal_mode != null);
+            testing.expectEqualStrings("wal", mem.spanZ(&journal_mode.?));
+        }
+
+        {
+            const journal_mode = try db.pragmaAlloc([]const u8, &arena.allocator, .{}, "journal_mode", arg);
+            testing.expect(journal_mode != null);
+            testing.expectEqualStrings("wal", journal_mode.?);
+        }
     }
 }
 
@@ -920,11 +983,9 @@ test "sqlite: read a single user into a struct" {
     var stmt = try db.prepare("SELECT id, name, age, weight FROM user WHERE id = ?{usize}");
     defer stmt.deinit();
 
-    var rows = try stmt.all(
-        TestUser,
-        .{ .allocator = &arena.allocator },
-        .{ .id = @as(usize, 20) },
-    );
+    var rows = try stmt.all(TestUser, &arena.allocator, .{}, .{
+        .id = @as(usize, 20),
+    });
     for (rows) |row| {
         testing.expectEqual(test_users[0].id, row.id);
         testing.expectEqualStrings(test_users[0].name, row.name);
@@ -933,14 +994,15 @@ test "sqlite: read a single user into a struct" {
 
     // Read a row with db.one()
     {
-        var row = try db.one(
+        var row = try db.oneAlloc(
             struct {
                 id: usize,
                 name: Text,
                 age: usize,
             },
+            &arena.allocator,
             "SELECT id, name, age FROM user WHERE id = ?{usize}",
-            .{ .allocator = &arena.allocator },
+            .{},
             .{@as(usize, 20)},
         );
         testing.expect(row != null);
@@ -963,11 +1025,7 @@ test "sqlite: read all users into a struct" {
     var stmt = try db.prepare("SELECT id, name, age, weight FROM user");
     defer stmt.deinit();
 
-    var rows = try stmt.all(
-        TestUser,
-        .{ .allocator = &arena.allocator },
-        .{},
-    );
+    var rows = try stmt.all(TestUser, &arena.allocator, .{}, .{});
     testing.expectEqual(@as(usize, 3), rows.len);
     for (rows) |row, i| {
         const exp = test_users[i];
@@ -988,7 +1046,7 @@ test "sqlite: read in an anonymous struct" {
     var stmt = try db.prepare("SELECT id, name, name, age, id, weight FROM user WHERE id = ?{usize}");
     defer stmt.deinit();
 
-    var row = try stmt.one(
+    var row = try stmt.oneAlloc(
         struct {
             id: usize,
             name: []const u8,
@@ -997,7 +1055,8 @@ test "sqlite: read in an anonymous struct" {
             is_id: bool,
             weight: f64,
         },
-        .{ .allocator = &arena.allocator },
+        &arena.allocator,
+        .{},
         .{ .id = @as(usize, 20) },
     );
     testing.expect(row != null);
@@ -1022,13 +1081,14 @@ test "sqlite: read in a Text struct" {
     var stmt = try db.prepare("SELECT id, name, age FROM user WHERE id = ?{usize}");
     defer stmt.deinit();
 
-    var row = try stmt.one(
+    var row = try stmt.oneAlloc(
         struct {
             id: usize,
             name: Text,
             age: usize,
         },
-        .{ .allocator = &arena.allocator },
+        &arena.allocator,
+        .{},
         .{@as(usize, 20)},
     );
     testing.expect(row != null);
@@ -1069,9 +1129,10 @@ test "sqlite: read a single text value" {
         var stmt: Statement(.{}, ParsedQuery.from(query)) = try db.prepare(query);
         defer stmt.deinit();
 
-        const name = try stmt.one(
+        const name = try stmt.oneAlloc(
             typ,
-            .{ .allocator = &arena.allocator },
+            &arena.allocator,
+            .{},
             .{ .id = @as(usize, 20) },
         );
         testing.expect(name != null);
@@ -1245,7 +1306,7 @@ test "sqlite: statement iterator" {
 
     var rows = std.ArrayList(Type).init(allocator);
     while (true) {
-        const row = (try iter.next(.{ .allocator = allocator })) orelse break;
+        const row = (try iter.nextAlloc(allocator, .{})) orelse break;
         try rows.append(row);
     }