dynamodb: add support for context management

This commit is contained in:
Muhammed Efe Cetin
2025-02-09 19:41:12 +03:00
parent 58bbe4632f
commit 9d7368cc48
2 changed files with 94 additions and 30 deletions

View File

@@ -17,9 +17,8 @@ import (
// Storage interface that is implemented by storage providers
type Storage struct {
db *awsdynamodb.Client
table string
requestTimeout time.Duration
db *awsdynamodb.Client
table string
}
// "k" is used as table column name for the key.
@@ -77,11 +76,8 @@ func New(config Config) *Storage {
return store
}
// Get value by key
func (s *Storage) Get(key string) ([]byte, error) {
ctx, cancel := s.requestContext()
defer cancel()
// GetWithContext value by key with context
func (s *Storage) GetWithContext(ctx context.Context, key string) ([]byte, error) {
k := make(map[string]types.AttributeValue)
k[keyAttrName] = &types.AttributeValueMemberS{
Value: key,
@@ -108,11 +104,12 @@ func (s *Storage) Get(key string) ([]byte, error) {
return item.V, err
}
// Set key with value
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
ctx, cancel := s.requestContext()
defer cancel()
func (s *Storage) Get(key string) ([]byte, error) {
return s.GetWithContext(context.Background(), key)
}
// Set key with value
func (s *Storage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
// Ain't Nobody Got Time For That
if len(key) <= 0 || len(val) <= 0 {
return nil
@@ -134,11 +131,12 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
return err
}
// Delete entry by key
func (s *Storage) Delete(key string) error {
ctx, cancel := s.requestContext()
defer cancel()
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
return s.SetWithContext(context.Background(), key, val, exp)
}
// Delete entry by key
func (s *Storage) DeleteWithContext(ctx context.Context, key string) error {
// Ain't Nobody Got Time For That
if len(key) <= 0 {
return nil
@@ -157,11 +155,12 @@ func (s *Storage) Delete(key string) error {
return err
}
// Reset all entries, including unexpired
func (s *Storage) Reset() error {
ctx, cancel := s.requestContext()
defer cancel()
func (s *Storage) Delete(key string) error {
return s.DeleteWithContext(context.Background(), key)
}
// Reset all entries, including unexpired
func (s *Storage) ResetWithContext(ctx context.Context) error {
deleteTableInput := awsdynamodb.DeleteTableInput{
TableName: &s.table,
}
@@ -169,14 +168,17 @@ func (s *Storage) Reset() error {
return err
}
func (s *Storage) Reset() error {
return s.ResetWithContext(context.Background())
}
// Close the database
func (s *Storage) Close() error {
return nil
}
func (s *Storage) createTable(cfg Config, describeTableInput awsdynamodb.DescribeTableInput) error {
ctx, cancel := s.requestContext()
defer cancel()
ctx := context.Background()
keyAttrType := "S" // For "string"
keyType := "HASH" // As opposed to "RANGE"
@@ -225,14 +227,6 @@ func (s *Storage) createTable(cfg Config, describeTableInput awsdynamodb.Describ
return nil
}
// Context for making requests will timeout if a non-zero timeout is configured
func (s *Storage) requestContext() (context.Context, context.CancelFunc) {
if s.requestTimeout > 0 {
return context.WithTimeout(context.Background(), s.requestTimeout)
}
return context.Background(), func() {}
}
func returnAWSConfig(cfg Config) (aws.Config, error) {
if cfg.Credentials != (Credentials{}) {
credentials := credentials.NewStaticCredentialsProvider(cfg.Credentials.AccessKey, cfg.Credentials.SecretAccessKey, "")