11package cassandra
22
33import (
4+ "context"
45 "errors"
56 "fmt"
67 "io"
@@ -52,7 +53,7 @@ type Cassandra struct {
5253 config * Config
5354}
5455
55- func WithInstance (session * gocql.Session , config * Config ) (database.Driver , error ) {
56+ func WithInstance (ctx context. Context , session * gocql.Session , config * Config ) (database.Driver , error ) {
5657 if config == nil {
5758 return nil , ErrNilConfig
5859 } else if len (config .KeyspaceName ) == 0 {
@@ -76,14 +77,14 @@ func WithInstance(session *gocql.Session, config *Config) (database.Driver, erro
7677 config : config ,
7778 }
7879
79- if err := c .ensureVersionTable (); err != nil {
80+ if err := c .ensureVersionTable (ctx ); err != nil {
8081 return nil , err
8182 }
8283
8384 return c , nil
8485}
8586
86- func (c * Cassandra ) Open (url string ) (database.Driver , error ) {
87+ func (c * Cassandra ) Open (ctx context. Context , url string ) (database.Driver , error ) {
8788 u , err := nurl .Parse (url )
8889 if err != nil {
8990 return nil , err
@@ -185,34 +186,34 @@ func (c *Cassandra) Open(url string) (database.Driver, error) {
185186 }
186187 }
187188
188- return WithInstance (session , & Config {
189+ return WithInstance (ctx , session , & Config {
189190 KeyspaceName : strings .TrimPrefix (u .Path , "/" ),
190191 MigrationsTable : u .Query ().Get ("x-migrations-table" ),
191192 MultiStatementEnabled : u .Query ().Get ("x-multi-statement" ) == "true" ,
192193 MultiStatementMaxSize : multiStatementMaxSize ,
193194 })
194195}
195196
196- func (c * Cassandra ) Close () error {
197+ func (c * Cassandra ) Close (ctx context. Context ) error {
197198 c .session .Close ()
198199 return nil
199200}
200201
201- func (c * Cassandra ) Lock () error {
202+ func (c * Cassandra ) Lock (ctx context. Context ) error {
202203 if ! c .isLocked .CAS (false , true ) {
203204 return database .ErrLocked
204205 }
205206 return nil
206207}
207208
208- func (c * Cassandra ) Unlock () error {
209+ func (c * Cassandra ) Unlock (ctx context. Context ) error {
209210 if ! c .isLocked .CAS (true , false ) {
210211 return database .ErrNotLocked
211212 }
212213 return nil
213214}
214215
215- func (c * Cassandra ) Run (migration io.Reader ) error {
216+ func (c * Cassandra ) Run (ctx context. Context , migration io.Reader ) error {
216217 if c .config .MultiStatementEnabled {
217218 var err error
218219 if e := multistmt .Parse (migration , multiStmtDelimiter , c .config .MultiStatementMaxSize , func (m []byte ) bool {
@@ -243,7 +244,7 @@ func (c *Cassandra) Run(migration io.Reader) error {
243244 return nil
244245}
245246
246- func (c * Cassandra ) SetVersion (version int , dirty bool ) error {
247+ func (c * Cassandra ) SetVersion (ctx context. Context , version int , dirty bool ) error {
247248 // DELETE instead of TRUNCATE because AWS Keyspaces does not support it
248249 // see: https://docs.aws.amazon.com/keyspaces/latest/devguide/cassandra-apis.html
249250 squery := `SELECT version FROM "` + c .config .MigrationsTable + `"`
@@ -273,7 +274,7 @@ func (c *Cassandra) SetVersion(version int, dirty bool) error {
273274}
274275
275276// Return current keyspace version
276- func (c * Cassandra ) Version () (version int , dirty bool , err error ) {
277+ func (c * Cassandra ) Version (ctx context. Context ) (version int , dirty bool , err error ) {
277278 query := `SELECT version, dirty FROM "` + c .config .MigrationsTable + `" LIMIT 1`
278279 err = c .session .Query (query ).Scan (& version , & dirty )
279280 switch {
@@ -291,7 +292,7 @@ func (c *Cassandra) Version() (version int, dirty bool, err error) {
291292 }
292293}
293294
294- func (c * Cassandra ) Drop () error {
295+ func (c * Cassandra ) Drop (ctx context. Context ) error {
295296 // select all tables in current schema
296297 query := fmt .Sprintf (`SELECT table_name from system_schema.tables WHERE keyspace_name='%s'` , c .config .KeyspaceName )
297298 iter := c .session .Query (query ).Iter ()
@@ -309,13 +310,13 @@ func (c *Cassandra) Drop() error {
309310// ensureVersionTable checks if versions table exists and, if not, creates it.
310311// Note that this function locks the database, which deviates from the usual
311312// convention of "caller locks" in the Cassandra type.
312- func (c * Cassandra ) ensureVersionTable () (err error ) {
313- if err = c .Lock (); err != nil {
313+ func (c * Cassandra ) ensureVersionTable (ctx context. Context ) (err error ) {
314+ if err = c .Lock (ctx ); err != nil {
314315 return err
315316 }
316317
317318 defer func () {
318- if e := c .Unlock (); e != nil {
319+ if e := c .Unlock (ctx ); e != nil {
319320 if err == nil {
320321 err = e
321322 } else {
@@ -328,7 +329,7 @@ func (c *Cassandra) ensureVersionTable() (err error) {
328329 if err != nil {
329330 return err
330331 }
331- if _ , _ , err = c .Version (); err != nil {
332+ if _ , _ , err = c .Version (ctx ); err != nil {
332333 return err
333334 }
334335 return nil
0 commit comments