balance and users list

cmd/balance.go

@@ -0,0 +1,35 @@
+package cmd
+
+import (
+	"flag"
+	"fmt"
+)
+
+func balance(args []string) error {
+	fs := flag.NewFlagSet("balance", flag.ExitOnError)
+	var userId int64
+	fs.Int64Var(&userId, "userId", 0, "Id of user (required)")
+	fs.Parse(args)
+
+	if userId <= 0 {
+		fs.Usage()
+		return fmt.Errorf("Missing userId")
+	}
+
+	tx, err := readDb()
+	if err != nil {
+		return err
+	}
+
+	total := float64(0)
+
+	for _, r := range tx.Records {
+		if r.UserId == uint64(userId) {
+			total += r.Amount
+		}
+	}
+	fmt.Printf("UserId: %d\n", userId)
+	fmt.Printf("Balance: %f\n", total)
+
+	return nil
+}

cmd/root.go

@@ -1,31 +1,30 @@
 package cmd
 
 import (
-	"flag"
 	"fmt"
-	"log"
 	"os"
 
 	"jxs.me/proto/mps"
 )
 
-func Root() error {
-	args := flag.NewFlagSet("root", flag.ExitOnError)
-	args.Parse(os.Args[1:])
-
-	source := args.Arg(0)
-	if source == "" {
-		return fmt.Errorf("File is required")
-	}
-
+func readDb() (*mps.TransactionLog, error) {
+	source := "txnlog.dat"
 	sourceFile, err := os.Open(source)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	defer sourceFile.Close()
-	tx, err := mps.Parse(sourceFile)
 
-	log.Printf("Parsed %d records", tx.Header.RecordCount)
+	return mps.Parse(sourceFile)
+}
 
-	return err
+func Root(args []string) error {
+	switch args[0] {
+	case "users":
+		return users(args[1:])
+	case "balance":
+		return balance(args[1:])
+	default:
+		return fmt.Errorf("Unknown command: %s", args[0])
+	}
 }

cmd/users.go

@@ -0,0 +1,24 @@
+package cmd
+
+import "fmt"
+
+type zilch struct{}
+
+func users(args []string) error {
+	tx, err := readDb()
+	if err != nil {
+		return err
+	}
+	idMap := map[uint64]zilch{}
+
+	for _, r := range tx.Records {
+		idMap[r.UserId] = zilch{}
+	}
+
+	fmt.Printf("Found %d users\n", len(idMap))
+	for id := range idMap {
+		fmt.Println(id)
+	}
+
+	return nil
+}

main.go

@@ -2,12 +2,13 @@ package main
 
 import (
 	"log"
+	"os"
 
 	"jxs.me/proto/cmd"
 )
 
 func main() {
-	if err := cmd.Root(); err != nil {
+	if err := cmd.Root(os.Args[1:]); err != nil {
 		log.Fatal(err)
 	}
 }