codecrafters submit [skip ci]

1 files changed, 37 insertions, 17 deletions

diff --git a/src/main.zig b/src/main.zig
index c06bd66..4394d13 100755
--- a/src/main.zig
+++ b/src/main.zig
@@ -38,23 +38,22 @@ pub fn main() !void {
     } else if (std.mem.eql(u8, args[2], ".tables")) {
         try tables.showTables(allocator, &file, page_size, stdout);
     } else if (std.mem.startsWith(u8, args[2], "SELECT") or std.mem.startsWith(u8, args[2], "select")) {
-        // Parse SELECT query: "SELECT column FROM table"
-        var tokens = std.mem.tokenizeScalar(u8, args[2], ' ');
-
-        // Skip "SELECT"
-        _ = tokens.next();
+        // Parse SELECT query
+        // Find "FROM" or "from" keyword
+        const from_idx_upper = std.mem.indexOf(u8, args[2], "FROM");
+        const from_idx_lower = std.mem.indexOf(u8, args[2], "from");
+        const from_idx = from_idx_upper orelse from_idx_lower orelse return error.InvalidQuery;
 
-        // Get column name or aggregate function
-        const column_name = tokens.next() orelse return error.InvalidQuery;
+        // Extract column part (between SELECT and FROM)
+        const select_len: usize = 6; // length of "SELECT" or "select"
+        const column_part = std.mem.trim(u8, args[2][select_len..from_idx], " \t\n");
 
-        // Skip "FROM" or "from"
-        _ = tokens.next();
-
-        // Get table name
-        const table_name = tokens.next() orelse return error.InvalidQuery;
+        // Extract table name (after FROM)
+        const from_end = from_idx + 4; // length of "FROM"
+        const table_name = std.mem.trim(u8, args[2][from_end..], " \t\n");
 
         // Check if this is COUNT(*)
-        if (std.mem.indexOf(u8, column_name, "count(") != null or std.mem.indexOf(u8, column_name, "COUNT(") != null) {
+        if (std.mem.indexOf(u8, column_part, "count(") != null or std.mem.indexOf(u8, column_part, "COUNT(") != null) {
             const rootpage = try schema.getRootpage(allocator, &file, page_size, table_name);
             const row_count = try schema.countRows(allocator, &file, page_size, rootpage);
             try stdout.print("{}\n", .{row_count});
@@ -63,14 +62,35 @@ pub fn main() !void {
             const create_sql = try schema.getCreateTableSQL(allocator, &file, page_size, table_name);
             defer allocator.free(create_sql);
 
-            // Find the column index
-            const column_idx = try schema.parseColumnIndex(create_sql, column_name);
+            // Parse column names (split by comma)
+            var column_list = std.ArrayList([]const u8){};
+            defer column_list.deinit(allocator);
+
+            var col_tokens = std.mem.tokenizeScalar(u8, column_part, ',');
+            while (col_tokens.next()) |col| {
+                const trimmed = std.mem.trim(u8, col, " \t\n");
+                try column_list.append(allocator, trimmed);
+            }
 
             // Get the table's root page
             const rootpage = try schema.getRootpage(allocator, &file, page_size, table_name);
 
-            // Read and print all rows
-            try schema.readTableRows(allocator, &file, page_size, rootpage, column_idx, stdout);
+            if (column_list.items.len == 1) {
+                // Single column query
+                const column_idx = try schema.parseColumnIndex(create_sql, column_list.items[0]);
+                try schema.readTableRows(allocator, &file, page_size, rootpage, column_idx, stdout);
+            } else {
+                // Multiple column query
+                var column_indices = std.ArrayList(usize){};
+                defer column_indices.deinit(allocator);
+
+                for (column_list.items) |col_name| {
+                    const idx = try schema.parseColumnIndex(create_sql, col_name);
+                    try column_indices.append(allocator, idx);
+                }
+
+                try schema.readTableRowsMultiColumn(allocator, &file, page_size, rootpage, column_indices.items, stdout);
+            }
         }
     } else {
         var tokens = std.mem.tokenizeScalar(u8, args[2], ' ');