Work on ResultIterator

2 files changed, 61 insertions, 22 deletions

diff --git a/src/conn/conn.zig b/src/conn/conn.zig
index a5a5ac6..fbde06c 100644
--- a/src/conn/conn.zig
+++ b/src/conn/conn.zig
@@ -6,6 +6,7 @@ const proto = @import("../proto/proto.zig");
 const StartupMessage = proto.StartupMessage;
 const PasswordMessage = proto.PasswordMessage;
 const BackendMessage = proto.BackendMessage;
+const RowDescription = proto.RowDescription;
 const read_message = proto.read_message;
 const ProtocolError = @import("../main.zig").ProtocolError;
 const ServerError = @import("../main.zig").ServerError;
@@ -52,7 +53,7 @@ pub fn connect(config: Config) !Conn {
     defer sm.deinit(allocator);
     try sm.write(allocator, writer);
     lp: while (true) {
-        var anymsg = try res.receiveMessage();
+        var anymsg = try res.receive_message();
         defer anymsg.deinit(allocator);
         switch (anymsg) {
             .ReadyForQuery => {
@@ -81,7 +82,7 @@ pub fn connect(config: Config) !Conn {
 
 // Messages should always be received through this function.
 // this'll handle generic stuff that should happen on the connection
-fn receiveMessage(self: *Conn) !BackendMessage {
+fn receive_message(self: *Conn) !BackendMessage {
     var anymsg = try read_message(self.allocator, self.stream.reader());
     errdefer anymsg.deinit(self.allocator);
     switch (anymsg) {
@@ -118,27 +119,64 @@ pub fn deinit(self: *Conn) void {
     self.stream.close();
 }
 
-// How to handle this ...
-// The Go code relies on polymorphism to generically read any message type.
-// I _could_ have a tagged union type thing
-// pub const ResultIterator = struct {
-//     conn: *Conn,
-//     command_concluded: bool = false,
-//     // NextRow advances the ResultReader to the next row and returns true if a row is available.
-//     pub fn next_row(self: *ResultIterator) bool {
-//         // TODO implement
-//         var reader = self.conn.stream.reader();
-//         switch (try reader.readByte()) {
-//             case
-//         }
-//         return false;
-//     }
-// };
+pub const ResultIterator = struct {
+    conn: *Conn,
+    multi_iterator: ?*MultiResultIterator = null,
+    row_description: ?proto.RowDescription = null,
+    current_datarow: ?proto.DataRow = null,
+    command_complete: ?proto.CommandComplete = null,
+    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 {
+        while (self.command_complete == null) {
+            var msg = try self.conn.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;
+                },
+                else => {
+                    msg.deinit(self.conn.allocator);
+                },
+            }
+        }
+        return false;
+    }
+    // row returns the current row data
+    pub fn row(self: ResultIterator) ?[][]const u8 {
+        return if (self.current_datarow) |dr| dr.columns else null;
+    }
+};
 
-// pub const MultiResultIterator = struct {
-//     conn: *Conn,
-//     fn next() ?
-// };
+pub const MultiResultIterator = struct {
+    conn: *Conn,
+    cri: ?*ResultIterator = null,
+
+    pub fn next_result(self: *MultiResultIterator) !bool {
+        if ()
+    }
+
+    pub fn result(self: MultiResultIterator) ?*ResultIterator {
+        return self.cri;
+    }
+    fn receive_message(self: *MultiResultIterator) !BackendMessage {
+        _ = self;
+        return error.NotImplemented;
+    }
+};
 
 // pub fn exec(self: *Conn) {
 
diff --git a/src/main.zig b/src/main.zig
index 9ce8de9..1c74915 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -11,6 +11,7 @@ pub const ProtocolError = error{
     WrongMessageType,
     InvalidTransactionStatus,
     InvalidFormatCode,
+    UnexpectedMessage,
 };
 
 pub const FormatCode = enum(u16) {