diff --git a/table/table.go b/table/table.go index 0ead82d..5b5e283 100644 --- a/table/table.go +++ b/table/table.go @@ -135,24 +135,32 @@ func (t *Table) SelectAll() (stmt string, names []string) { return qb.Select(t.metadata.Name).ToCql() } -// Insert returns insert all columns statement. -func (t *Table) Insert() (stmt string, names []string) { - return t.insert.stmt, t.insert.names +// Insert returns insert statement. +func (t *Table) Insert(columns ...string) (stmt string, names []string) { + if len(columns) == 0 { + return t.insert.stmt, t.insert.names + } + + return t.InsertBuilder(columns...).ToCql() } -// InsertQuery returns query which inserts all columns. -func (t *Table) InsertQuery(session gocqlx.Session) *gocqlx.Queryx { - return session.Query(t.Insert()) +// InsertQuery returns query which inserts columns. +func (t *Table) InsertQuery(session gocqlx.Session, columns ...string) *gocqlx.Queryx { + return session.Query(t.Insert(columns...)) } -// InsertQueryContext returns query wrapped with context which inserts all columns. -func (t *Table) InsertQueryContext(ctx context.Context, session gocqlx.Session) *gocqlx.Queryx { - return t.InsertQuery(session).WithContext(ctx) +// InsertQueryContext returns query wrapped with context which inserts columns. +func (t *Table) InsertQueryContext(ctx context.Context, session gocqlx.Session, columns ...string) *gocqlx.Queryx { + return t.InsertQuery(session, columns...).WithContext(ctx) } -// InsertBuilder returns a builder initialised with all columns. -func (t *Table) InsertBuilder() *qb.InsertBuilder { - return qb.Insert(t.metadata.Name).Columns(t.metadata.Columns...) +// InsertBuilder returns a builder initialised with columns. +func (t *Table) InsertBuilder(columns ...string) *qb.InsertBuilder { + if len(columns) == 0 { + return qb.Insert(t.metadata.Name).Columns(t.metadata.Columns...) + } + + return qb.Insert(t.metadata.Name).Columns(columns...) } // Update returns update by primary key statement. diff --git a/table/table_test.go b/table/table_test.go index 32871de..509ecae 100644 --- a/table/table_test.go +++ b/table/table_test.go @@ -116,6 +116,7 @@ func TestTableSelect(t *testing.T) { func TestTableInsert(t *testing.T) { table := []struct { M Metadata + C []string N []string S string }{ @@ -129,10 +130,32 @@ func TestTableInsert(t *testing.T) { N: []string{"a", "b", "c", "d"}, S: "INSERT INTO table (a,b,c,d) VALUES (?,?,?,?) ", }, + { + M: Metadata{ + Name: "table", + Columns: []string{"a", "b", "c", "d"}, + PartKey: []string{"a"}, + SortKey: []string{"b"}, + }, + C: []string{"a", "b", "c"}, + N: []string{"a", "b", "c"}, + S: "INSERT INTO table (a,b,c) VALUES (?,?,?) ", + }, + } + + for _, test := range table { + stmt, names := New(test.M).Insert(test.C...) + if diff := cmp.Diff(test.S, stmt); diff != "" { + t.Error(diff) + } + if diff := cmp.Diff(test.N, names); diff != "" { + t.Error(diff, names) + } } + // run InsertBuilder on the same data set for _, test := range table { - stmt, names := New(test.M).Insert() + stmt, names := New(test.M).InsertBuilder(test.C...).ToCql() if diff := cmp.Diff(test.S, stmt); diff != "" { t.Error(diff) }