From cc68867ad0d67fcb8070cf2217ff4e49bee059b0 Mon Sep 17 00:00:00 2001
From: Yemin Li <44004590+YeminLi@users.noreply.github.com>
Date: Mon, 30 Sep 2024 16:25:14 -0400
Subject: [PATCH] expose SelectBuilder for table Get operation (#251)

* expose select builder for get

* add comment
---
 table/table.go      |  9 +++++++++
 table/table_test.go | 11 +++++++++++
 2 files changed, 20 insertions(+)

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) {