qb: select

qb/expr.go

@@ -15,14 +15,13 @@ type expr interface {
 
 type columns []string
 
-func (cols columns) writeCql(cql *bytes.Buffer) (names []string) {
+func (cols columns) writeCql(cql *bytes.Buffer) {
 	for i, c := range cols {
 		cql.WriteString(c)
 		if i < len(cols)-1 {
 			cql.WriteByte(',')
 		}
 	}
-	return
 }
 
 type using struct {
@@ -30,7 +29,7 @@ type using struct {
 	ttl       time.Duration
 }
 
-func (u using) writeCql(cql *bytes.Buffer) (names []string) {
+func (u using) writeCql(cql *bytes.Buffer) {
 	var v []string
 	if !u.timestamp.IsZero() {
 		v = append(v, fmt.Sprint("TIMESTAMP ", u.timestamp.UnixNano()/1000))
@@ -43,8 +42,6 @@ func (u using) writeCql(cql *bytes.Buffer) (names []string) {
 		cql.WriteString(strings.Join(v, ","))
 		cql.WriteByte(' ')
 	}
-
-	return
 }
 
 type where []expr

qb/select.go

@@ -0,0 +1,153 @@
+package qb
+
+// SELECT reference:
+// http://docs.datastax.com/en/dse/5.1/cql/cql/cql_reference/cql_commands/cqlSelect.html
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+)
+
+type Order bool
+
+const (
+	ASC  Order = true
+	DESC       = false
+)
+
+type SelectBuilder struct {
+	table             string
+	columns           columns
+	distinct          columns
+	where             where
+	groupBy           columns
+	orderBy           string
+	order             Order
+	limit             uint
+	limitPerPartition uint
+	allowFiltering    bool
+}
+
+// Select returns a new SelectBuilder with the given table name.
+func Select(table string) *SelectBuilder {
+	return &SelectBuilder{
+		table: table,
+	}
+}
+
+func (b *SelectBuilder) ToCql() (stmt string, names []string, err error) {
+	if b.table == "" {
+		err = errors.New("select statements must specify a table")
+		return
+	}
+	if b.distinct != "" && len(b.columns) > 0 {
+		err = fmt.Errorf("select statements must specify either a column list or DISTINCT partition_key")
+		return
+	}
+
+	cql := bytes.Buffer{}
+
+	cql.WriteString("SELECT ")
+	switch {
+	case len(b.distinct) > 0:
+		cql.WriteString("DISTINCT ")
+		b.distinct.writeCql(&cql)
+	case len(b.groupBy) > 0:
+		b.groupBy.writeCql(&cql)
+		cql.WriteByte(',')
+		b.columns.writeCql(&cql)
+	case len(b.columns) == 0:
+		cql.WriteByte('*')
+	default:
+		b.columns.writeCql(&cql)
+	}
+	cql.WriteString(" FROM ")
+	cql.WriteString(b.table)
+	cql.WriteByte(' ')
+
+	names = b.where.writeCql(&cql)
+
+	if len(b.groupBy) > 0 {
+		cql.WriteString("GROUP BY ")
+		b.groupBy.writeCql(&cql)
+		cql.WriteByte(' ')
+	}
+
+	if b.orderBy != "" {
+		cql.WriteString("ORDER BY ")
+		cql.WriteString(b.orderBy)
+		if b.order {
+			cql.WriteString(" ASC ")
+		} else {
+			cql.WriteString(" DESC ")
+		}
+	}
+
+	if b.limit != 0 {
+		cql.WriteString("LIMIT ")
+		cql.WriteString(fmt.Sprint(b.limit))
+		cql.WriteByte(' ')
+	}
+
+	if b.limitPerPartition != 0 {
+		cql.WriteString("PER PARTITION LIMIT ")
+		cql.WriteString(fmt.Sprint(b.limitPerPartition))
+		cql.WriteByte(' ')
+	}
+
+	if b.allowFiltering {
+		cql.WriteString("ALLOW FILTERING ")
+	}
+
+	stmt = cql.String()
+	return
+}
+
+// From sets the table to be selected from.
+func (b *SelectBuilder) From(table string) *SelectBuilder {
+	b.table = table
+	return b
+}
+
+func (b *SelectBuilder) Columns(columns ...string) *SelectBuilder {
+	b.columns = append(b.columns, columns...)
+	return b
+}
+
+func (b *SelectBuilder) Distinct(columns... string) *SelectBuilder {
+	b.distinct = append(b.distinct, columns...)
+	return b
+}
+
+func (b *SelectBuilder) Where(e ...expr) *SelectBuilder {
+	b.where = append(b.where, e...)
+	return b
+}
+
+// GroupBy sets GROUP BY clause on the query. Columns must be a primary key,
+// this will automatically add the the columns as first selectors.
+func (b *SelectBuilder) GroupBy(columns... string) *SelectBuilder {
+	b.groupBy = append(b.groupBy, columns...)
+	return b
+}
+
+func (b *SelectBuilder) OrderBy(column string, o Order) *SelectBuilder {
+	b.orderBy, b.order = column, o
+	return b
+}
+
+func (b *SelectBuilder) Limit(limit uint) *SelectBuilder {
+	b.limit = limit
+	return b
+}
+
+func (b *SelectBuilder) LimitPerPartition(limit uint) *SelectBuilder {
+	b.limitPerPartition = limit
+	return b
+}
+
+func (b *SelectBuilder) AllowFiltering() *SelectBuilder {
+	b.allowFiltering = true
+	return b
+}

qb/select_test.go

@@ -0,0 +1,106 @@
+package qb
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestSelectBuilder(t *testing.T) {
+	m := mockExpr{
+		cql:   "expr",
+		names: []string{"expr"},
+	}
+
+	table := []struct {
+		B *SelectBuilder
+		N []string
+		S string
+	}{
+		// Basic test for select *
+		{
+			B: Select("cycling.cyclist_name"),
+			S: "SELECT * FROM cycling.cyclist_name ",
+		},
+		// Basic test for select columns
+		{
+			B: Select("cycling.cyclist_name").Columns("id", "user_uuid", "firstname"),
+			S: "SELECT id,user_uuid,firstname FROM cycling.cyclist_name ",
+		},
+		// Basic test for select distinct
+		{
+			B: Select("cycling.cyclist_name").Distinct("id"),
+			S: "SELECT DISTINCT id FROM cycling.cyclist_name ",
+		},
+		// Change table name
+		{
+			B: Select("cycling.cyclist_name").From("Foobar"),
+			S: "SELECT * FROM Foobar ",
+		},
+		// Add WHERE
+		{
+			B: Select("cycling.cyclist_name").Where(m).Where(mockExpr{
+				cql:   "expr_1",
+				names: []string{"expr_1"},
+			}, mockExpr{
+				cql:   "expr_2",
+				names: []string{"expr_2"},
+			}),
+			S: "SELECT * FROM cycling.cyclist_name WHERE expr AND expr_1 AND expr_2 ",
+			N: []string{"expr", "expr_1", "expr_2"},
+		},
+		// Add GROUP BY
+		{
+			B: Select("cycling.cyclist_name").Columns("MAX(stars) as max_stars").GroupBy("id"),
+			S: "SELECT id,MAX(stars) as max_stars FROM cycling.cyclist_name GROUP BY id ",
+		},
+		// Add ORDER BY
+		{
+			B: Select("cycling.cyclist_name").Where(m).OrderBy("firstname", ASC),
+			S: "SELECT * FROM cycling.cyclist_name WHERE expr ORDER BY firstname ASC ",
+			N: []string{"expr"},
+		},
+		// Add LIMIT
+		{
+			B: Select("cycling.cyclist_name").Where(m).Limit(10),
+			S: "SELECT * FROM cycling.cyclist_name WHERE expr LIMIT 10 ",
+			N: []string{"expr"},
+		},
+		// Add PER PARTITION LIMIT
+		{
+			B: Select("cycling.cyclist_name").Where(m).LimitPerPartition(10),
+			S: "SELECT * FROM cycling.cyclist_name WHERE expr PER PARTITION LIMIT 10 ",
+			N: []string{"expr"},
+		},
+		// Add ALLOW FILTERING
+		{
+			B: Select("cycling.cyclist_name").Where(m).AllowFiltering(),
+			S: "SELECT * FROM cycling.cyclist_name WHERE expr ALLOW FILTERING ",
+			N: []string{"expr"},
+		},
+	}
+
+	for _, test := range table {
+		stmt, names, err := test.B.ToCql()
+		if err != nil {
+			t.Error("unexpected error", err)
+		}
+		if diff := cmp.Diff(test.S, stmt); diff != "" {
+			t.Error(diff)
+		}
+		if diff := cmp.Diff(test.N, names); diff != "" {
+			t.Error(diff)
+		}
+	}
+}
+
+func BenchmarkSelectBuilder(b *testing.B) {
+	m := mockExpr{
+		cql:   "expr",
+		names: []string{"expr"},
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		Select("cycling.cyclist_name").Columns("id", "user_uuid", "firstname", "stars").Where(m)
+	}
+}