diff --git a/cmd/bricksllm/main.go b/cmd/bricksllm/main.go index a95a273..035e08d 100644 --- a/cmd/bricksllm/main.go +++ b/cmd/bricksllm/main.go @@ -162,6 +162,11 @@ func main() { log.Sugar().Fatalf("error creating user id for users table: %v", err) } + err = store.InitializeSecondaryKeyTable() + if err != nil { + log.Sugar().Fatalf("error initializing secondary key table: %v", err) + } + go store.PrepareEventsIndexes(log) cpMemStore, err := memdb.NewCustomProvidersMemDb(store, log, cfg.InMemoryDbUpdateInterval) @@ -190,7 +195,7 @@ func main() { rateLimitRedisCache := redis.NewClient(defaultRedisOption(cfg, 0)) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := rateLimitRedisCache.Ping(ctx).Err(); err != nil { + if err = rateLimitRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to rate limit redis cache: %v", err) } @@ -198,7 +203,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := costLimitRedisCache.Ping(ctx).Err(); err != nil { + if err = costLimitRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to cost limit redis cache: %v", err) } @@ -206,7 +211,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := costRedisStorage.Ping(ctx).Err(); err != nil { + if err = costRedisStorage.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to cost limit redis storage: %v", err) } @@ -214,7 +219,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := apiRedisCache.Ping(ctx).Err(); err != nil { + if err = apiRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to api redis cache: %v", err) } @@ -222,7 +227,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := accessRedisCache.Ping(ctx).Err(); err != nil { + if err = accessRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to api redis cache: %v", err) } @@ -230,7 +235,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := userRateLimitRedisCache.Ping(ctx).Err(); err != nil { + if err = userRateLimitRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to user rate limit redis cache: %v", err) } @@ -238,7 +243,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := userCostLimitRedisCache.Ping(ctx).Err(); err != nil { + if err = userCostLimitRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to user cost limit redis cache: %v", err) } @@ -246,7 +251,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := userCostRedisStorage.Ping(ctx).Err(); err != nil { + if err = userCostRedisStorage.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to user cost redis cache: %v", err) } @@ -254,7 +259,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := userAccessRedisCache.Ping(ctx).Err(); err != nil { + if err = userAccessRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to user access redis storage: %v", err) } @@ -262,7 +267,7 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := providerSettingsRedisCache.Ping(ctx).Err(); err != nil { + if err = providerSettingsRedisCache.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to provider settings redis storage: %v", err) } @@ -278,10 +283,18 @@ func main() { ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if err := requestsLimitRedisStorage.Ping(ctx).Err(); err != nil { + if err = requestsLimitRedisStorage.Ping(ctx).Err(); err != nil { log.Sugar().Fatalf("error connecting to requests limit redis storage: %v", err) } + secondaryKeysRedisCache := redis.NewClient(defaultRedisOption(cfg, 12)) + + ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err = secondaryKeysRedisCache.Ping(ctx).Err(); err != nil { + log.Sugar().Fatalf("error connecting to secondary keys redis storage: %v", err) + } + rateLimitCache := redisStorage.NewCache(rateLimitRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) costLimitCache := redisStorage.NewCache(costLimitRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) costStorage := redisStorage.NewStore(costRedisStorage, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) @@ -295,6 +308,7 @@ func main() { psCache := redisStorage.NewProviderSettingsCache(providerSettingsRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) keysCache := redisStorage.NewKeysCache(keysRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) + secondaryKeysCache := redisStorage.NewSecondaryKeysCache(secondaryKeysRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) requestsLimitStorage := redisStorage.NewStore(requestsLimitRedisStorage, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) encryptor, err := encryptor.NewEncryptor(cfg.DecryptionEndpoint, cfg.EncryptionEndpoint, cfg.EnableEncrytion, cfg.EncryptionTimeout, cfg.Audience) @@ -303,7 +317,7 @@ func main() { } v := validator.NewValidator(costLimitCache, rateLimitCache, costStorage, requestsLimitStorage) - m := manager.NewManager(store, costLimitCache, rateLimitCache, accessCache, keysCache, requestsLimitStorage) + m := manager.NewManager(store, costLimitCache, rateLimitCache, accessCache, keysCache, secondaryKeysCache, requestsLimitStorage) krm := manager.NewReportingManager(costStorage, store, store, v) psm := manager.NewProviderSettingsManager(store, psCache, encryptor) cpm := manager.NewCustomProvidersManager(store, cpMemStore) diff --git a/internal/authenticator/authenticator.go b/internal/authenticator/authenticator.go index 115211f..f846ab3 100644 --- a/internal/authenticator/authenticator.go +++ b/internal/authenticator/authenticator.go @@ -31,6 +31,7 @@ type routesManager interface { type keysCache interface { GetKeyViaCache(hash string) (*key.ResponseKey, error) + GetKeyHashBySecondary(kHash string) (string, error) } type keyStorage interface { @@ -206,6 +207,19 @@ func anonymize(input string) string { return string(input[0:5]) + "**********************************************" } +const secondaryPrefix = "secondary_" + +func (a *Authenticator) getHashViaSecondary(rawKey string) (string, error) { + if len(rawKey) > 0 && rawKey[0] != secondaryPrefix[0] { + return hasher.Hash(rawKey), nil + } + if !strings.HasPrefix(rawKey, secondaryPrefix) { + return hasher.Hash(rawKey), nil + } + hash := hasher.Hash(rawKey) + return a.kc.GetKeyHashBySecondary(hash) +} + func (a *Authenticator) AuthenticateHttpRequest(req *http.Request, xCustomProviderId string) (*key.ResponseKey, []*provider.Setting, error) { var raw string var err error @@ -224,19 +238,20 @@ func (a *Authenticator) AuthenticateHttpRequest(req *http.Request, xCustomProvid return nil, nil, err } - hash := hasher.Hash(raw) + hash, err := a.getHashViaSecondary(raw) - key, err := a.kc.GetKeyViaCache(hash) - if key != nil { + rKey, err := a.kc.GetKeyViaCache(hash) + if rKey != nil { telemetry.Incr(metricname.COUNTER_AUTHENTICATOR_FOUND_KEY_FROM_MEMDB, nil, 1) } - if key == nil { - key, err = a.kc.GetKeyViaCache(raw) + if rKey == nil { + rKey, err = a.kc.GetKeyViaCache(raw) } if err != nil { - _, ok := err.(notFoundError) + var nFoundError notFoundError + ok := errors.As(err, &nFoundError) if ok { return nil, nil, internal_errors.NewAuthError(fmt.Sprintf("key %s is not found", anonymize(raw))) } @@ -244,11 +259,11 @@ func (a *Authenticator) AuthenticateHttpRequest(req *http.Request, xCustomProvid return nil, nil, err } - if key == nil { + if rKey == nil { return nil, nil, internal_errors.NewAuthError(fmt.Sprintf("key %s is not found", anonymize(raw))) } - if key.Revoked { + if rKey.Revoked { return nil, nil, internal_errors.NewAuthError(fmt.Sprintf("key %s has been revoked", anonymize(raw))) } @@ -271,17 +286,17 @@ func (a *Authenticator) AuthenticateHttpRequest(req *http.Request, xCustomProvid default: return nil, nil, errors.New("invalid xCustomAuth location") } - return key, settings, nil + return rKey, settings, nil } if strings.HasPrefix(req.URL.Path, "/api/routes") { - err = a.canKeyAccessCustomRoute(req.URL.Path, key.KeyId) + err = a.canKeyAccessCustomRoute(req.URL.Path, rKey.KeyId) if err != nil { return nil, nil, err } } - settingIds := key.GetSettingIds() + settingIds := rKey.GetSettingIds() allSettings := []*provider.Setting{} selected := []*provider.Setting{} for _, settingId := range settingIds { @@ -308,7 +323,7 @@ func (a *Authenticator) AuthenticateHttpRequest(req *http.Request, xCustomProvid if len(selected) != 0 { used := selected[0] - if key.RotationEnabled { + if rKey.RotationEnabled { used = selected[rand.Intn(len(selected))] } @@ -337,7 +352,7 @@ func (a *Authenticator) AuthenticateHttpRequest(req *http.Request, xCustomProvid return nil, nil, err } - return key, selected, nil + return rKey, selected, nil } return nil, nil, internal_errors.NewAuthError(fmt.Sprintf("provider setting not found for key %s", raw)) diff --git a/internal/manager/key.go b/internal/manager/key.go index f4f5b2d..08c445e 100644 --- a/internal/manager/key.go +++ b/internal/manager/key.go @@ -11,6 +11,7 @@ import ( "github.com/bricks-cloud/bricksllm/internal/key" "github.com/bricks-cloud/bricksllm/internal/policy" "github.com/bricks-cloud/bricksllm/internal/provider" + secondarykey "github.com/bricks-cloud/bricksllm/internal/secondary-key" "github.com/bricks-cloud/bricksllm/internal/telemetry" "github.com/bricks-cloud/bricksllm/internal/util" ) @@ -26,6 +27,9 @@ type Storage interface { GetProviderSettings(withSecret bool, ids []string) ([]*provider.Setting, error) GetKey(keyId string) (*key.ResponseKey, error) GetKeyByHash(hash string) (*key.ResponseKey, error) + GetKeyHashBySecondary(sHash string) (string, error) + CreateSecondaryKey(secondaryHash string) error + UpdateSecondaryKey(secondaryHash, keyHash string) error } type costLimitCache interface { @@ -46,27 +50,35 @@ type keyCache interface { Get(keyId string) (*key.ResponseKey, error) } +type secondaryKeyCache interface { + Set(sHash string, value string, ttl time.Duration) error + Delete(sHash string) error + Get(sHash string) (string, error) +} + type requestsLimitStorage interface { DeleteCounter(keyId string) error } type Manager struct { - s Storage - clc costLimitCache - rlc rateLimitCache - ac accessCache - kc keyCache - rqls requestsLimitStorage + s Storage + clc costLimitCache + rlc rateLimitCache + ac accessCache + kc keyCache + secondaryKC secondaryKeyCache + rqls requestsLimitStorage } -func NewManager(s Storage, clc costLimitCache, rlc rateLimitCache, ac accessCache, kc keyCache, rqls requestsLimitStorage) *Manager { +func NewManager(s Storage, clc costLimitCache, rlc rateLimitCache, ac accessCache, kc keyCache, secondaryKC secondaryKeyCache, rqls requestsLimitStorage) *Manager { return &Manager{ - s: s, - clc: clc, - rlc: rlc, - ac: ac, - kc: kc, - rqls: rqls, + s: s, + clc: clc, + rlc: rlc, + ac: ac, + kc: kc, + secondaryKC: secondaryKC, + rqls: rqls, } } @@ -241,6 +253,59 @@ func (m *Manager) GetKeyViaCache(raw string) (*key.ResponseKey, error) { return k, nil } +func (m *Manager) GetKeyHashBySecondary(sHash string) (string, error) { + h, _ := m.secondaryKC.Get(sHash) + if h == "" { + telemetry.Incr("bricksllm.manager.get_key_hash_by_secondary.cache_miss", nil, 1) + stored, err := m.s.GetKeyHashBySecondary(sHash) + if err != nil { + return "", err + } + if stored == "" { + return "", errors.New("key hash not found") + } + err = m.secondaryKC.Set(sHash, stored, 24*time.Hour) + if err != nil { + telemetry.Incr("bricksllm.manager.get_key_hash_by_secondary.set_error", nil, 1) + } + h = stored + } + telemetry.Incr("bricksllm.manager.get_key_hash_by_secondary.cache_hit", nil, 1) + return h, nil +} + +func (m *Manager) CreateSecondaryKey(keyCreate secondarykey.SecondaryKeyCreate) error { + if keyCreate.Key == "" { + return errors.New("key is required for creating secondary key") + } + return m.s.CreateSecondaryKey(hasher.Hash(keyCreate.Key)) +} + +func (m *Manager) UpdateSecondaryKey(keyUpdate secondarykey.SecondaryKeyUpdate) error { + if keyUpdate.Key == "" { + return errors.New("key is required for updating secondary key") + } + if keyUpdate.LinkedKeyId == "" { + return errors.New("linkedKeyId is required for updating secondary key") + } + rKey, err := m.s.GetKey(keyUpdate.LinkedKeyId) + if err != nil { + return err + } + if rKey == nil { + return errors.New("linked key not found for updating secondary key") + } + err = m.s.UpdateSecondaryKey(hasher.Hash(keyUpdate.Key), rKey.Key) + if err != nil { + return err + } + err = m.secondaryKC.Set(hasher.Hash(keyUpdate.Key), rKey.Key, 24*time.Hour) + if err != nil { + telemetry.Incr("bricksllm.manager.update_secondary_key.set_cache_error", nil, 1) + } + return nil +} + func (m *Manager) DeleteKey(id string) error { return m.s.DeleteKey(id) } diff --git a/internal/secondary-key/secondary_key.go b/internal/secondary-key/secondary_key.go new file mode 100644 index 0000000..378af24 --- /dev/null +++ b/internal/secondary-key/secondary_key.go @@ -0,0 +1,10 @@ +package secondary_key + +type SecondaryKeyCreate struct { + Key string `json:"key"` +} + +type SecondaryKeyUpdate struct { + Key string `json:"key"` + LinkedKeyId string `json:"linkedKeyId"` +} diff --git a/internal/server/web/admin/admin.go b/internal/server/web/admin/admin.go index 4aa928b..d9fcd63 100644 --- a/internal/server/web/admin/admin.go +++ b/internal/server/web/admin/admin.go @@ -13,6 +13,7 @@ import ( "github.com/bricks-cloud/bricksllm/internal/policy" "github.com/bricks-cloud/bricksllm/internal/provider" "github.com/bricks-cloud/bricksllm/internal/provider/custom" + secondarykey "github.com/bricks-cloud/bricksllm/internal/secondary-key" "github.com/bricks-cloud/bricksllm/internal/telemetry" "github.com/bricks-cloud/bricksllm/internal/util" "github.com/gin-gonic/gin" @@ -31,6 +32,8 @@ type KeyManager interface { GetKeysV2(tags, keyIds []string, revoked *bool, limit, offset int, name, order string, returnCount bool) (*key.GetKeysResponse, error) UpdateKey(id string, key *key.UpdateKey) (*key.ResponseKey, error) CreateKey(key *key.RequestKey) (*key.ResponseKey, error) + CreateSecondaryKey(keyCreate secondarykey.SecondaryKeyCreate) error + UpdateSecondaryKey(keyUpdate secondarykey.SecondaryKeyUpdate) error DeleteKey(id string) error } @@ -85,6 +88,9 @@ func NewAdminServer(log *zap.Logger, mode string, m KeyManager, krm KeyReporting router.PATCH("/api/key-management/keys/:id", getUpdateKeyHandler(m, prod)) router.DELETE("/api/key-management/keys/:id", getDeleteKeyHandler(m, prod)) + router.POST("/api/v2/key-management/secondary-keys", getCreateSecondaryKeyHandler(m, prod)) + router.PATCH("/api/v2/key-management/secondary-keys/:id", getUpdateSecondaryKeyHandler(m, prod)) + router.GET("/api/reporting/keys/:id", getGetKeyReportingHandler(krm, prod)) router.POST("/api/reporting/events", getGetEventMetricsHandler(krm, prod)) router.POST("/api/reporting/events-by-day", getGetEventMetricsByDayHandler(krm, prod)) @@ -567,6 +573,166 @@ func getCreateKeyHandler(m KeyManager, prod bool) gin.HandlerFunc { } } +func getCreateSecondaryKeyHandler(m KeyManager, prod bool) gin.HandlerFunc { + return func(c *gin.Context) { + log := util.GetLogFromCtx(c) + telemetry.Incr("bricksllm.admin.get_create_secondary_key_handler.requests", nil, 1) + + start := time.Now() + defer func() { + dur := time.Since(start) + telemetry.Timing("bricksllm.admin.get_create_secondary_key_handler.latency", dur, nil, 1) + }() + + //path := "/api/key-management/keys" + path := c.FullPath() + // TODO: DEBUG + fmt.Println("------------- full path: ", c.FullPath()) + + if c == nil || c.Request == nil { + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/empty-context", + Title: "context is empty error", + Status: http.StatusInternalServerError, + Detail: "gin context is empty", + Instance: path, + }) + return + } + + data, err := io.ReadAll(c.Request.Body) + if err != nil { + logError(log, "error when reading secondary key creation request body", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/request-body-read", + Title: "request body reader error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + secondaryKeyCreate := &secondarykey.SecondaryKeyCreate{} + err = json.Unmarshal(data, secondaryKeyCreate) + if err != nil { + logError(log, "error when unmarshalling secondary key creation request body", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/json-unmarshal", + Title: "json unmarshaller error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + err = m.CreateSecondaryKey(*secondaryKeyCreate) + if err != nil { + errType := "internal" + + defer func() { + telemetry.Incr("bricksllm.admin.get_create_secondary_key_handler.create_secondary_key_error", []string{ + "error_type:" + errType, + }, 1) + }() + + logError(log, "error when creating a secondary key", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/key-manager", + Title: "secondary key creation error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + telemetry.Incr("bricksllm.admin.get_create_secondary_key_handler.success", nil, 1) + + c.Status(http.StatusOK) + } +} + +func getUpdateSecondaryKeyHandler(m KeyManager, prod bool) gin.HandlerFunc { + return func(c *gin.Context) { + log := util.GetLogFromCtx(c) + telemetry.Incr("bricksllm.admin.get_update_secondary_key_handler.requests", nil, 1) + + start := time.Now() + defer func() { + dur := time.Since(start) + telemetry.Timing("bricksllm.admin.get_update_secondary_key_handler.latency", dur, nil, 1) + }() + + //path := "/api/key-management/keys" + path := c.FullPath() + // TODO: DEBUG + fmt.Println("------------- full path: ", c.FullPath()) + + if c == nil || c.Request == nil { + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/empty-context", + Title: "context is empty error", + Status: http.StatusInternalServerError, + Detail: "gin context is empty", + Instance: path, + }) + return + } + + data, err := io.ReadAll(c.Request.Body) + if err != nil { + logError(log, "error when reading secondary key update request body", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/request-body-read", + Title: "request body reader error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + secondaryKeyUpdate := &secondarykey.SecondaryKeyUpdate{} + err = json.Unmarshal(data, secondaryKeyUpdate) + if err != nil { + logError(log, "error when unmarshalling secondary key creation request body", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/json-unmarshal", + Title: "json unmarshaller error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + err = m.UpdateSecondaryKey(*secondaryKeyUpdate) + if err != nil { + errType := "internal" + + defer func() { + telemetry.Incr("bricksllm.admin.get_update_secondary_key_handler.update_secondary_key_error", []string{ + "error_type:" + errType, + }, 1) + }() + + logError(log, "error when updating a secondary key", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/key-manager", + Title: "secondary key updating error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + telemetry.Incr("bricksllm.admin.get_update_secondary_key_handler.success", nil, 1) + + c.Status(http.StatusOK) + } +} + func getUpdateProviderSettingHandler(m ProviderSettingsManager, prod bool) gin.HandlerFunc { return func(c *gin.Context) { log := util.GetLogFromCtx(c) diff --git a/internal/storage/postgresql/secondary_key.go b/internal/storage/postgresql/secondary_key.go new file mode 100644 index 0000000..43162d5 --- /dev/null +++ b/internal/storage/postgresql/secondary_key.go @@ -0,0 +1,57 @@ +package postgresql + +import ( + "context" +) + +func (s *Store) InitializeSecondaryKeyTable() error { + createTableQuery := ` + CREATE TABLE IF NOT EXISTS secondary_keys ( + secondary_hash VARCHAR(255) PRIMARY KEY, + key_hash VARCHAR(255), + )` + + ctxTimeout, cancel := context.WithTimeout(context.Background(), s.wt) + defer cancel() + _, err := s.db.ExecContext(ctxTimeout, createTableQuery) + if err != nil { + return err + } + return nil +} + +func (s *Store) GetKeyHashBySecondary(sHash string) (string, error) { + ctxTimeout, cancel := context.WithTimeout(context.Background(), s.rt) + defer cancel() + query := "SELECT key_hash FROM secondary_keys WHERE secondary_hash = $1" + + row := s.db.QueryRowContext(ctxTimeout, query, sHash) + + var keyHash string + err := row.Scan(&keyHash) + if err != nil { + return "", err + } + + return keyHash, nil +} + +func (s *Store) CreateSecondaryKey(secondaryHash string) error { + ctxTimeout, cancel := context.WithTimeout(context.Background(), s.wt) + defer cancel() + + query := "INSERT INTO secondary_keys (secondary_hash) VALUES ($1)" + + _, err := s.db.ExecContext(ctxTimeout, query, secondaryHash) + return err +} + +func (s *Store) UpdateSecondaryKey(secondaryHash, keyHash string) error { + ctxTimeout, cancel := context.WithTimeout(context.Background(), s.rt) + defer cancel() + + query := "UPDATE secondary_keys SET key_hash = $1 WHERE secondary_hash = $2" + + _, err := s.db.ExecContext(ctxTimeout, query, keyHash, secondaryHash) + return err +} diff --git a/internal/storage/redis/seconadry-key-cache.go b/internal/storage/redis/seconadry-key-cache.go new file mode 100644 index 0000000..9203a23 --- /dev/null +++ b/internal/storage/redis/seconadry-key-cache.go @@ -0,0 +1,62 @@ +package redis + +import ( + "context" + "time" + + "github.com/redis/go-redis/v9" +) + +type SecondaryKeysCache struct { + client *redis.Client + wt time.Duration + rt time.Duration +} + +func NewSecondaryKeysCache(c *redis.Client, wt time.Duration, rt time.Duration) *SecondaryKeysCache { + return &SecondaryKeysCache{ + client: c, + wt: wt, + rt: rt, + } +} + +func (c *SecondaryKeysCache) Set(sHash string, value string, ttl time.Duration) error { + ctx, cancel := context.WithTimeout(context.Background(), c.wt) + defer cancel() + err := c.client.Set(ctx, sHash, value, ttl).Err() + if err != nil { + return err + } + + return nil +} + +func (c *SecondaryKeysCache) Delete(sHash string) error { + ctx, cancel := context.WithTimeout(context.Background(), c.wt) + defer cancel() + err := c.client.Del(ctx, sHash).Err() + if err != nil { + return err + } + + return nil +} + +func (c *SecondaryKeysCache) Get(sHash string) (string, error) { + ctx, cancel := context.WithTimeout(context.Background(), c.rt) + defer cancel() + + result := c.client.Get(ctx, sHash) + err := result.Err() + if err != nil { + return "", err + } + + bs, err := result.Bytes() + if err != nil { + return "", err + } + + return string(bs), nil +}