Added Context variants of Query functions

scylladb · Jun 19, 2020 · 4ea6f42 · 4ea6f42
1 parent 8af6506
commit 4ea6f42
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 0 deletions.
diff --git a/qb/batch.go b/qb/batch.go
@@ -6,6 +6,7 @@ package qb
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"time"
 
@@ -71,6 +72,11 @@ func (b *BatchBuilder) Query(session gocqlx.Session) *gocqlx.Queryx {
 	return session.Query(b.ToCql())
 }
 
+// QueryContext returns query wrapped with context built on top of current BatchBuilder state.
+func (b *BatchBuilder) QueryContext(ctx context.Context, session gocqlx.Session) *gocqlx.Queryx {
+	return b.Query(session).WithContext(ctx)
+}
+
 // Add builds the builder and adds the statement to the batch.
 func (b *BatchBuilder) Add(builder Builder) *BatchBuilder {
 	return b.AddStmt(builder.ToCql())

diff --git a/qb/delete.go b/qb/delete.go
@@ -9,6 +9,7 @@ package qb
 
 import (
 	"bytes"
+	"context"
 	"time"
 
 	"github.com/scylladb/gocqlx/v2"
@@ -61,6 +62,11 @@ func (b *DeleteBuilder) Query(session gocqlx.Session) *gocqlx.Queryx {
 	return session.Query(b.ToCql())
 }
 
+// QueryContext returns query wrapped with context built on top of current DeleteBuilder state.
+func (b *DeleteBuilder) QueryContext(ctx context.Context, session gocqlx.Session) *gocqlx.Queryx {
+	return b.Query(session).WithContext(ctx)
+}
+
 // From sets the table to be deleted from.
 func (b *DeleteBuilder) From(table string) *DeleteBuilder {
 	b.table = table

diff --git a/qb/insert.go b/qb/insert.go
@@ -9,6 +9,7 @@ package qb
 
 import (
 	"bytes"
+	"context"
 	"time"
 
 	"github.com/scylladb/gocqlx/v2"
@@ -85,6 +86,11 @@ func (b *InsertBuilder) Query(session gocqlx.Session) *gocqlx.Queryx {
 	return session.Query(b.ToCql())
 }
 
+// QueryContext returns query wrapped with context built on top of current InsertBuilder state.
+func (b *InsertBuilder) QueryContext(ctx context.Context, session gocqlx.Session) *gocqlx.Queryx {
+	return b.Query(session).WithContext(ctx)
+}
+
 // Into sets the INTO clause of the query.
 func (b *InsertBuilder) Into(table string) *InsertBuilder {
 	b.table = table

diff --git a/qb/select.go b/qb/select.go
@@ -9,6 +9,7 @@ package qb
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 
 	"github.com/scylladb/gocqlx/v2"
@@ -125,6 +126,11 @@ func (b *SelectBuilder) Query(session gocqlx.Session) *gocqlx.Queryx {
 	return session.Query(b.ToCql())
 }
 
+// QueryContext returns query wrapped with context built on top of current SelectBuilder state.
+func (b *SelectBuilder) QueryContext(ctx context.Context, session gocqlx.Session) *gocqlx.Queryx {
+	return b.Query(session).WithContext(ctx)
+}
+
 // From sets the table to be selected from.
 func (b *SelectBuilder) From(table string) *SelectBuilder {
 	b.table = table

diff --git a/qb/update.go b/qb/update.go
@@ -9,6 +9,7 @@ package qb
 
 import (
 	"bytes"
+	"context"
 	"time"
 
 	"github.com/scylladb/gocqlx/v2"
@@ -80,6 +81,11 @@ func (b *UpdateBuilder) Query(session gocqlx.Session) *gocqlx.Queryx {
 	return session.Query(b.ToCql())
 }
 
+// QueryContext returns query wrapped with context built on top of current UpdateBuilder state.
+func (b *UpdateBuilder) QueryContext(ctx context.Context, session gocqlx.Session) *gocqlx.Queryx {
+	return b.Query(session).WithContext(ctx)
+}
+
 // Table sets the table to be updated.
 func (b *UpdateBuilder) Table(table string) *UpdateBuilder {
 	b.table = table

diff --git a/table/table.go b/table/table.go
@@ -5,6 +5,8 @@
 package table
 
 import (
+	"context"
+
 	"github.com/scylladb/gocqlx/v2"
 	"github.com/scylladb/gocqlx/v2/qb"
 )
@@ -95,6 +97,11 @@ func (t *Table) GetQuery(session gocqlx.Session, columns ...string) *gocqlx.Quer
 	return session.Query(t.Get(columns...))
 }
 
+// GetQueryContext returns query wrapped with context which gets by partition key.
+func (t *Table) GetQueryContext(ctx context.Context, session gocqlx.Session, columns ...string) *gocqlx.Queryx {
+	return t.GetQuery(session, columns...).WithContext(ctx)
+}
+
 // Select returns select by partition key statement.
 func (t *Table) Select(columns ...string) (stmt string, names []string) {
 	if len(columns) == 0 {
@@ -112,6 +119,11 @@ func (t *Table) SelectQuery(session gocqlx.Session, columns ...string) *gocqlx.Q
 	return session.Query(t.Select(columns...))
 }
 
+// SelectQueryContext returns query wrapped with context which selects by partition key statement.
+func (t *Table) SelectQueryContext(ctx context.Context, session gocqlx.Session, columns ...string) *gocqlx.Queryx {
+	return t.SelectQuery(session, columns...).WithContext(ctx)
+}
+
 // SelectBuilder returns a builder initialised to select by partition key
 // statement.
 func (t *Table) SelectBuilder(columns ...string) *qb.SelectBuilder {
@@ -128,6 +140,11 @@ func (t *Table) InsertQuery(session gocqlx.Session) *gocqlx.Queryx {
 	return session.Query(t.Insert())
 }
 
+// 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)
+}
+
 // Update returns update by primary key statement.
 func (t *Table) Update(columns ...string) (stmt string, names []string) {
 	return t.UpdateBuilder(columns...).ToCql()
@@ -138,6 +155,11 @@ func (t *Table) UpdateQuery(session gocqlx.Session, columns ...string) *gocqlx.Q
 	return session.Query(t.Update(columns...))
 }
 
+// UpdateQueryContext returns query wrapped with context which updates by primary key.
+func (t *Table) UpdateQueryContext(ctx context.Context, session gocqlx.Session, columns ...string) *gocqlx.Queryx {
+	return t.UpdateQuery(session, columns...).WithContext(ctx)
+}
+
 // UpdateBuilder returns a builder initialised to update by primary key statement.
 func (t *Table) UpdateBuilder(columns ...string) *qb.UpdateBuilder {
 	return qb.Update(t.metadata.Name).Set(columns...).Where(t.primaryKeyCmp...)
@@ -153,6 +175,11 @@ func (t *Table) DeleteQuery(session gocqlx.Session, columns ...string) *gocqlx.Q
 	return session.Query(t.Delete(columns...))
 }
 
+// DeleteQueryContext returns query wrapped with context which delete by primary key.
+func (t *Table) DeleteQueryContext(ctx context.Context, session gocqlx.Session, columns ...string) *gocqlx.Queryx {
+	return t.DeleteQuery(session, columns...).WithContext(ctx)
+}
+
 // DeleteBuilder returns a builder initialised to delete by primary key statement.
 func (t *Table) DeleteBuilder(columns ...string) *qb.DeleteBuilder {
 	return qb.Delete(t.metadata.Name).Columns(columns...).Where(t.primaryKeyCmp...)