Consolidate all code to main, refactor

main.go

@@ -2,11 +2,179 @@ package main
 
 import (
 	"bytes"
+	"fmt"
+	"os"
+	"strings"
+
+	"container/list"
+	"database/sql"
+
 	"github.com/babashka/pod-babashka-sqlite3/babashka"
-	"github.com/babashka/pod-babashka-sqlite3/pod"
+	_ "github.com/mattn/go-sqlite3" // Import go-sqlite3 library
 	"github.com/russolsen/transit"
 )
 
+func debug(v interface{}) {
+	fmt.Fprintf(os.Stderr, "debug: %+v\n", v)
+}
+
+func encodeRows(rows *sql.Rows) ([]interface{}, error) {
+	cols, err := rows.Columns()
+	columns := make([]transit.Keyword, len(cols))
+	for i, col := range cols {
+		columns[i] = transit.Keyword(col)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	var data []interface{}
+
+	values := make([]interface{}, len(columns))
+	scanArgs := make([]interface{}, len(values))
+	for i := range values {
+		scanArgs[i] = &values[i]
+	}
+
+	for rows.Next() {
+		results := make(map[transit.Keyword]interface{})
+
+		if err = rows.Scan(scanArgs...); err != nil {
+			debug(err)
+			return nil, err
+		}
+
+		for i, val := range values {
+			col := columns[i]
+			results[col] = val
+		}
+
+		// debug(results)
+		// debug(fmt.Sprintf("%T", results))
+
+		data = append(data, results)
+	}
+
+	return data, nil
+}
+
+type ExecResult = map[transit.Keyword]int64
+
+func encodeResult(result sql.Result) (ExecResult, error) {
+	rowsAffected, err := result.RowsAffected()
+	lastInsertedId, err := result.LastInsertId()
+
+	if err != nil {
+		return nil, err
+	}
+
+	res := ExecResult{
+		transit.Keyword("rows-affected"):    rowsAffected,
+		transit.Keyword("last-inserted-id"): lastInsertedId,
+	}
+	return res, nil
+}
+
+func listToSlice(l *list.List) []interface{} {
+	slice := make([]interface{}, l.Len())
+	cnt := 0
+	for e := l.Front(); e != nil; e = e.Next() {
+		slice[cnt] = e.Value
+		cnt++
+	}
+
+	return slice
+}
+
+func parseQuery(args string) (string, string, []interface{}, error) {
+	reader := strings.NewReader(args)
+	decoder := transit.NewDecoder(reader)
+	value, err := decoder.Decode()
+	if err != nil {
+		return "", "", nil, err
+	}
+
+	argSlice := listToSlice(value.(*list.List))
+	db := argSlice[0].(string)
+	queryArgs := argSlice[1].([]interface{})
+	query := queryArgs[0].(string)
+
+	return db, query, queryArgs[1:], nil
+}
+
+func makeArgs(query []string) []interface{} {
+	args := make([]interface{}, len(query)-1)
+
+	for i := range query[1:] {
+		args[i] = query[i+1]
+	}
+
+	return args
+}
+
+func processMessage(message *babashka.Message) (interface{}, error) {
+	switch message.Op {
+	case "describe":
+		return &babashka.DescribeResponse{
+			Format: "transit+json",
+			Namespaces: []babashka.Namespace{
+				{
+					Name: "pod.babashka.sqlite3",
+					Vars: []babashka.Var{
+						{
+							Name: "execute!",
+						},
+						{
+							Name: "query!",
+						},
+					},
+				},
+			},
+		}, nil
+	case "invoke":
+		db, query, args, err := parseQuery(message.Args)
+		if err != nil {
+			return nil, err
+		}
+
+		conn, err := sql.Open("sqlite3", db)
+		if err != nil {
+			return nil, err
+		}
+
+		defer conn.Close()
+
+		switch message.Var {
+		case "pod.babashka.sqlite3/execute!":
+			res, err := conn.Exec(query, args...)
+			if err != nil {
+				return nil, err
+			}
+
+			if json, err := encodeResult(res); err != nil {
+				return nil, err
+			} else {
+				return json, nil
+			}
+		case "pod.babashka.sqlite3/query!":
+			res, err := conn.Query(query, args...)
+			if err != nil {
+				return nil, err
+			}
+
+			if json, err := encodeRows(res); err != nil {
+				return nil, err
+			} else {
+				return json, nil
+			}
+		default:
+			return nil, fmt.Errorf("Unknown var %s", message.Var)
+		}
+	default:
+		return nil, fmt.Errorf("Unknown op %s", message.Op)
+	}
+}
+
 func main() {
 	for {
 		message, err := babashka.ReadMessage()
@@ -15,7 +183,7 @@ func main() {
 			continue
 		}
 
-		res, err := pod.ProcessMessage(message)
+		res, err := processMessage(message)
 		if err != nil {
 			babashka.WriteErrorResponse(message, err)
 			continue
@@ -32,7 +200,6 @@ func main() {
 		if err := encoder.Encode(res); err != nil {
 			babashka.WriteErrorResponse(message, err)
 		} else {
-			//println("buf", buf.String())
 			babashka.WriteInvokeResponse(message, string(buf.String()))
 		}
 	}