1 files changed, 45 insertions, 27 deletions

diff --git a/src/conn/conn.zig b/src/conn/conn.zig
index fbde06c..db5084c 100644
--- a/src/conn/conn.zig
+++ b/src/conn/conn.zig
@@ -125,56 +125,74 @@ pub const ResultIterator = struct {
     row_description: ?proto.RowDescription = null,
     current_datarow: ?proto.DataRow = null,
     command_complete: ?proto.CommandComplete = null,
+    pub fn init(conn: *Conn) ResultIterator {
+        return .{
+            .conn = conn,
+        };
+    }
     pub fn deinit(self: *ResultIterator) void {
         if (self.row_description != null) self.row_description.?.deinit(self.conn.allocator);
         if (self.current_datarow != null) self.current_datarow.?.deinit(self.conn.allocator);
         if (self.command_complete != null) self.command_complete.?.deinit(self.conn.allocator);
     }
-    // NextRow advances the ResultIterator to the next row and returns true if a row is available.
-    pub fn next_row(self: *ResultIterator) !bool {
+    // NextRow advances the ResultIterator to the next row and returns a row if one is available.
+    // or null if we've reached the end of the reuslt.
+    pub fn next_row(self: *ResultIterator) !?[][]const u8 {
         while (self.command_complete == null) {
-            var msg = try self.conn.receive_message();
+            var msg = try self.receive_message();
             switch (msg) {
-                .DataRow => |dr| {
-                    if (self.row_description != null) self.row_description.?.deinit(self.conn.allocator);
-                    self.current_datarow = dr;
-                    return true;
-                },
-                .RowDescription => |rd| {
-                    if (self.row_description != null) return ProtocolError.UnexpectedMessage;
-                    self.row_description = rd;
-                },
-                .CommandComplete => |cc| {
-                    if (self.command_complete != null) return ProtocolError.UnexpectedMessage;
-                    self.command_complete = cc;
+                .DataRow => {
+                    return self.current_datarow.?.columns;
                 },
                 else => {
                     msg.deinit(self.conn.allocator);
                 },
             }
         }
-        return false;
+        return null;
+    }
+
+    pub fn skip_to_end(self: *ResultIterator) !void {
+        while (self.command_complete == null) {
+            _ = try self.receive_message();
+        }
     }
-    // row returns the current row data
-    pub fn row(self: ResultIterator) ?[][]const u8 {
-        return if (self.current_datarow) |dr| dr.columns else null;
+
+    fn receive_message(self: *ResultIterator) !BackendMessage {
+        var msg = if (self.multi_iterator == null) try self.conn.receive_message() else try self.multi_iterator.?.receive_message();
+        switch (msg) {
+            .DataRow => |dr| {
+                if (self.current_datarow != null) self.current_datarow.?.deinit(self.conn.allocator);
+                self.current_datarow = try dr.clone(self.conn.allocator);
+            },
+            .RowDescription => |rd| {
+                if (self.row_description != null) return ProtocolError.UnexpectedMessage;
+                self.row_description = try rd.clone(self.conn.allocator);
+            },
+            .CommandComplete => |cc| {
+                if (self.command_complete != null) return ProtocolError.UnexpectedMessage;
+                self.command_complete = try cc.clone(self.conn.allocator);
+            },
+        }
+        return msg;
     }
 };
 
 pub const MultiResultIterator = struct {
     conn: *Conn,
-    cri: ?*ResultIterator = null,
+    cri: ?*ResultIterator,
 
-    pub fn next_result(self: *MultiResultIterator) !bool {
-        if ()
+    // returns the next result iterator, or null if we've reached the end of the results
+    pub fn next_result(self: *MultiResultIterator) !?*ResultIterator {
+        if (self.cri != null) {
+            try self.cri.?.skip_to_end();
+        }
     }
 
-    pub fn result(self: MultiResultIterator) ?*ResultIterator {
-        return self.cri;
-    }
     fn receive_message(self: *MultiResultIterator) !BackendMessage {
-        _ = self;
-        return error.NotImplemented;
+        var msg = try self.conn.receive_message();
+        switch (msg) {}
+        return msg;
     }
 };