diff --git a/qb/batch.go b/qb/batch.go index 973b31a..7642fb9 100644 --- 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 index 089b8dc..d0d0070 100644 --- 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 index 0f8b7df..fd11bea 100644 --- 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 index 12d4249..3a9aff2 100644 --- 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 index fa3f8f8..b9d63b7 100644 --- 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 index 72dddfe..ae1b449 100644 --- 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...)