- handle query 'Select (col_name) from table_name'

app/helper.go

@@ -1,20 +1,34 @@
 package main
 
-func parseVarIntBtes(data []byte) (uint64, int) {
+import (
+	"regexp"
+	"strings"
+)
+
+func decodeVarint(data *[]byte, offset int64) (uint64, int) {
 	var result uint64
+	var i int64
 
-	offset := 0
+	for {
+		currentByte := (*data)[offset+i]
 
-	for offset < len(data) {
-		currentByte := data[offset]
-		result |= uint64(currentByte&0x7F) << (8 * (offset))
-		offset++
-		if currentByte&0x80 == 0 {
+		if i == 8 {
+			result <<= 8
+			result |= uint64(currentByte)
 			break
+		} else {
+			result <<= 7
 		}
-	}
 
-	return result, offset
+		result |= uint64(currentByte & 0x7f)
+
+		if currentByte>>7 == 0 {
+			break
+		}
+
+		i++
+	}
+	return result, int(i + 1)
 }
 
 func getContentSizeFromSerialType(serialType uint64) uint64 {
@@ -34,29 +48,14 @@ func getContentSizeFromSerialType(serialType uint64) uint64 {
 	}
 }
 
-func parseCellHeader(data []byte) ([]uint64, int) {
-	offset := 0
-
-	// Size of record header (varint) including it self:
-	// any byte that comes after offset + headerLength is the actual content
-	headerLength, size := parseVarIntBtes(data[offset:])
-
-	offset += size
-
-	var columnSizes []uint64
-
-	for offset < int(headerLength) {
-		// 1. get serial type
-		columnType, size := parseVarIntBtes(data[offset:])
-
-		// 2. get size
-		columnSize := getContentSizeFromSerialType(columnType)
-
-		columnSizes = append(columnSizes, columnSize)
-
-		offset += size
+func getColumnIndex(createStatement string, columnName string) int {
+	re := regexp.MustCompile(`CREATE TABLE \w+\s*\(([^\)]+)\)`)
+	matches := re.FindStringSubmatch(createStatement)
+	columns := strings.Split(matches[1], ",")
+	for i, c := range columns {
+		if strings.Split(strings.TrimSpace(c), " ")[0] == columnName {
+			return i
+		}
 	}
-
-	return columnSizes, int(headerLength)
-
+	return -1
 }

app/main.go

@@ -5,8 +5,9 @@ import (
 	"log"
 	"os"
 	"strings"
+
 	// Available if you need it!
-	// "github.com/xwb1989/sqlparser"
+	"github.com/xwb1989/sqlparser"
 )
 
 // Usage: your_program.sh sample.db .dbinfo
@@ -95,12 +96,27 @@ func main() {
 			log.Fatal(err)
 		}
 
-		command := strings.Split(command, " ")
-		tableName := command[len(command)-1]
+		handleQuery(command, databaseFile)
+
+	}
+}
+
+func handleQuery(queries string, databaseFile *os.File) {
+	parsedQuery, err := sqlparser.Parse(queries)
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	switch parsedQuery := parsedQuery.(type) {
+
+	case *sqlparser.Select:
+		tableName := sqlparser.String(parsedQuery.From[0])
+		selectExp := sqlparser.String(parsedQuery.SelectExprs[0])
 
 		header := make([]byte, 100)
 
-		_, err = databaseFile.Read(header)
+		_, err := databaseFile.Read(header)
 
 		if err != nil {
 			log.Fatal(err)
@@ -118,18 +134,40 @@ func main() {
 			log.Fatal(err)
 		}
 
-		// rootPages := make(map[string]uint8)
-		for _, row := range firstPage.Rows {
-			if string(row.Columns[0]) == "table" && string(row.Columns[1]) == tableName {
-				targetPage, err := readPage(databaseFile, int(row.Columns[3][0]), int(databaseHeader.PageSize))
+		if strings.Contains(strings.ToLower(selectExp), "count") {
 
-				if err != nil {
-					log.Fatal(err)
+			// rootPages := make(map[string]uint8)
+			for _, row := range firstPage.Rows {
+				if string(row.Columns[0]) == "table" && string(row.Columns[1]) == tableName {
+					targetPageHeader, err := peakPageHeader(databaseFile, int(row.Columns[3][0]), int(databaseHeader.PageSize))
+
+					if err != nil {
+						log.Fatal(err)
+					}
+
+					fmt.Println(targetPageHeader.CellCount)
+				}
+			}
+
+		} else {
+
+			for _, row := range firstPage.Rows {
+				if string(row.Columns[0]) == "table" && string(row.Columns[1]) == tableName {
+					columnIndex := getColumnIndex(string(row.Columns[4]), selectExp)
+					targetPage, err := readPage(databaseFile, int(row.Columns[3][0]), int(databaseHeader.PageSize))
+
+					if err != nil {
+						log.Fatal(err)
+					}
+
+					for _, row := range targetPage.Rows {
+						fmt.Println(string(row.Columns[columnIndex]))
+					}
 				}
 
-				fmt.Println(targetPage.Header.CellCount)
-				break
 			}
+
 		}
 	}
+
 }

app/page.go

@@ -3,7 +3,6 @@ package main
 import (
 	"encoding/binary"
 	"fmt"
-	"log"
 	"os"
 )
 
@@ -68,9 +67,38 @@ func parsePointers(data []byte) []uint16 {
 
 }
 
+func peakPageHeader(file *os.File, pageNumber int, pageSize int) (PageHeader, error) {
+	buff := make([]byte, pageSize)
+
+	readerOffset := pageSize * (pageNumber - 1)
+
+	_, err := file.ReadAt(buff, int64(readerOffset))
+
+	if err != nil {
+		return PageHeader{}, err
+	}
+
+	offset := 0
+
+	if pageNumber == 1 {
+		offset = 100
+	}
+
+	header, err := unmarshalPageHeader(buff[offset : offset+8])
+
+	if err != nil {
+		return PageHeader{}, err
+	}
+
+	return header, nil
+
+}
+
 func readPage(file *os.File, pageNumber int, pageSize int) (Page, error) {
 	// info, _ := file.Stat() // Size in bytes
 	// fmt.Printf("The size of the file is %d bytes.\n", info.Size())
+	// fmt.Println("target page index is ", pageNumber)
+	// fmt.Println("page size is ", pageSize)
 
 	buff := make([]byte, pageSize)
 
@@ -89,9 +117,7 @@ func readPage(file *os.File, pageNumber int, pageSize int) (Page, error) {
 	}
 
 	header, err := unmarshalPageHeader(buff[offset : offset+8])
-	if pageNumber != 1 {
-		fmt.Println(header.CellCount)
-	}
+
 	offset += 8
 
 	if err != nil {
@@ -106,60 +132,48 @@ func readPage(file *os.File, pageNumber int, pageSize int) (Page, error) {
 	var cells []Cell
 
 	for _, pointer := range pointers {
-		cell := readCell(file, int(pointer))
+		cell := readCell(buff, int(pointer))
 		cells = append(cells, cell)
 	}
 
 	return Page{Header: header, Rows: cells}, nil
 }
 
-func readCell(file *os.File, readerOffset int) Cell {
-
-	varintBytes := make([]byte, 8)
-
-	if _, err := file.ReadAt(varintBytes, int64(readerOffset)); err != nil {
-		log.Fatal(err)
-	}
-
-	// Size of the record
-	rowLength, size := parseVarIntBtes(varintBytes)
-
-	readerOffset += size
+func readCell(data []byte, offset int) Cell {
 
-	buff := make([]byte, rowLength-uint64(size))
-
-	if _, err := file.ReadAt(buff, int64(readerOffset)); err != nil {
-		log.Fatal(err)
-	}
+	rowLength, size := decodeVarint(&data, int64(offset))
 
-	offset := 0
+	offset += size
 
-	// The rowid (safe to ignore)
-	rowID, size := parseVarIntBtes(buff[offset:])
+	rowID, size := decodeVarint(&data, int64(offset))
 
 	offset += size
 
-	columnSizes, size := parseCellHeader(buff[offset:])
-
+	headerLength, size := decodeVarint(&data, int64(offset))
+	headerByteEnd := offset + int(headerLength)
 	offset += size
 
-	var body [][]byte
+	var columnSizes []uint64
 
-	for index, columnSize := range columnSizes {
-		body = append(body, buff[offset:offset+int(columnSize)])
+	for offset < headerByteEnd {
+		serialType, bytesRead := decodeVarint(&data, int64(offset))
+		columnSize := getContentSizeFromSerialType(serialType)
+		columnSizes = append(columnSizes, columnSize)
+		offset += bytesRead
+	}
 
-		if index == 3 {
-			break
-		}
+	var columns [][]byte
 
+	for _, columnSize := range columnSizes {
+		content := data[offset : offset+int(columnSize)]
+		columns = append(columns, content)
 		offset += int(columnSize)
 	}
 
 	return Cell{
-		Size:        rowLength,
-		RowID:       rowID,
-		HeaderSize:  uint64(size),
-		ColumnSizes: columnSizes,
-		Columns:     body,
+		Size:       rowLength,
+		RowID:      rowID,
+		HeaderSize: uint64(size),
+		Columns:    columns,
 	}
 }

go.mod

@@ -10,4 +10,4 @@ module github/com/codecrafters-io/sqlite-starter-go
 
 go 1.22
 
-require github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2
+require github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2