diff --git a/table/table.go b/table/table.go index e006ca8..b210852 100644 --- a/table/table.go +++ b/table/table.go @@ -102,6 +102,15 @@ func (t *Table) GetQueryContext(ctx context.Context, session gocqlx.Session, col return t.GetQuery(session, columns...).WithContext(ctx) } +// GetBuilder returns a builder initialised to select by primary key +func (t *Table) GetBuilder(columns ...string) *qb.SelectBuilder { + if len(columns) == 0 { + return qb.Select(t.metadata.Name).Where(t.primaryKeyCmp...) + } + + return qb.Select(t.metadata.Name).Columns(columns...).Where(t.primaryKeyCmp...) +} + // Select returns select by partition key statement. func (t *Table) Select(columns ...string) (stmt string, names []string) { if len(columns) == 0 { diff --git a/table/table_test.go b/table/table_test.go index da8464e..8911f92 100644 --- a/table/table_test.go +++ b/table/table_test.go @@ -60,6 +60,17 @@ func TestTableGet(t *testing.T) { t.Error(diff, names) } } + + // run GetBuilder on the same data set + for _, test := range table { + stmt, names := New(test.M).GetBuilder(test.C...).ToCql() + if diff := cmp.Diff(test.S, stmt); diff != "" { + t.Error(diff) + } + if diff := cmp.Diff(test.N, names); diff != "" { + t.Error(diff, names) + } + } } func TestTableSelect(t *testing.T) {