diff --git a/mssql/mssql.go b/mssql/mssql.go index ff7c5af5..3e332973 100644 --- a/mssql/mssql.go +++ b/mssql/mssql.go @@ -1,6 +1,7 @@ package mssql import ( + "context" "database/sql" "fmt" "net/url" @@ -132,13 +133,13 @@ func New(config ...Config) *Storage { return store } -// Get value by key -func (s *Storage) Get(key string) ([]byte, error) { +// GetWithContext gets value by key with context +func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) { if len(key) <= 0 { return nil, nil } - row := s.db.QueryRow(s.sqlSelect, key) + row := s.db.QueryRowContext(ctx, s.sqlSelect, key) var ( data = []byte{} @@ -161,8 +162,13 @@ func (s *Storage) Get(key string) ([]byte, error) { return data, nil } -// Set key with value -func (s *Storage) Set(key string, val []byte, exp time.Duration) error { +// Get gets value by key +func (s *Storage) Get(key string) ([]byte, error) { + return s.GetWithContext(context.Background(), key) +} + +// SetWithContext key with value and expiration time with context +func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error { if len(key) <= 0 || len(val) <= 0 { return nil } @@ -172,24 +178,39 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error { expSeconds = time.Now().Add(exp).Unix() } - _, err := s.db.Exec(s.sqlInsert, key, val, expSeconds) + _, err := s.db.ExecContext(ctx, s.sqlInsert, key, val, expSeconds) + return err +} + +// Set key with value and expiration time +func (s *Storage) Set(key string, val []byte, exp time.Duration) error { + return s.SetWithContext(context.Background(), key, val, exp) +} + +// DeleteWithContext entry by key with context +func (s *Storage) DeleteWithContext(ctx context.Context, key string) error { + if len(key) <= 0 { + return nil + } + + _, err := s.db.ExecContext(ctx, s.sqlDelete, key) return err } // Delete entry by key func (s *Storage) Delete(key string) error { - if len(key) <= 0 { - return nil - } + return s.DeleteWithContext(context.Background(), key) +} - _, err := s.db.Exec(s.sqlDelete, key) +// ResetWithContext all entries, including unexpired with context +func (s *Storage) ResetWithContext(ctx context.Context) error { + _, err := s.db.ExecContext(ctx, s.sqlReset) return err } // Reset all entries, including unexpired func (s *Storage) Reset() error { - _, err := s.db.Exec(s.sqlReset) - return err + return s.ResetWithContext(context.Background()) } // Close the database diff --git a/mssql/mssql_test.go b/mssql/mssql_test.go index bb209caf..501eb393 100644 --- a/mssql/mssql_test.go +++ b/mssql/mssql_test.go @@ -1,6 +1,7 @@ package mssql import ( + "context" "database/sql" "os" "testing" @@ -26,6 +27,19 @@ func Test_MSSQL_Set(t *testing.T) { require.NoError(t, err) } +func Test_MSSQL_SetWithContext(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := testStore.SetWithContext(ctx, key, val, 0) + require.ErrorIs(t, err, context.Canceled) +} + func Test_MSSQL_Set_Override(t *testing.T) { var ( key = "john" @@ -53,6 +67,23 @@ func Test_MSSQL_Get(t *testing.T) { require.Equal(t, val, result) } +func Test_MSSQL_GetWithContext(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + err := testStore.Set(key, val, 0) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result, err := testStore.GetWithContext(ctx, key) + require.ErrorIs(t, err, context.Canceled) + require.Zero(t, len(result)) +} + func Test_MSSQL_Set_Expiration(t *testing.T) { var ( key = "john" @@ -97,6 +128,26 @@ func Test_MSSQL_Delete(t *testing.T) { require.Zero(t, len(result)) } +func Test_MSSQL_DeleteWithContext(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + err := testStore.Set(key, val, 0) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = testStore.DeleteWithContext(ctx, key) + require.ErrorIs(t, err, context.Canceled) + + result, err := testStore.Get(key) + require.NoError(t, err) + require.Equal(t, val, result) +} + func Test_MSSQL_Reset(t *testing.T) { val := []byte("doe") @@ -118,6 +169,30 @@ func Test_MSSQL_Reset(t *testing.T) { require.Zero(t, len(result)) } +func Test_MSSQL_ResetWithContext(t *testing.T) { + val := []byte("doe") + + err := testStore.Set("john1", val, 0) + require.NoError(t, err) + + err = testStore.Set("john2", val, 0) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = testStore.ResetWithContext(ctx) + require.ErrorIs(t, err, context.Canceled) + + result, err := testStore.Get("john1") + require.NoError(t, err) + require.Equal(t, val, result) + + result, err = testStore.Get("john2") + require.NoError(t, err) + require.Equal(t, val, result) +} + func Test_MSSQL_GC(t *testing.T) { testVal := []byte("doe")