/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ package plugin import ( "context" "fmt" "math/rand/v2" "time" "github.com/apache/answer/internal/entity" "github.com/segmentfault/pacman/cache" "github.com/segmentfault/pacman/log" "xorm.io/builder" "xorm.io/xorm" ) // Error variables for KV storage operations var ( // ErrKVKeyNotFound is returned when the requested key does not exist in the KV storage ErrKVKeyNotFound = fmt.Errorf("key not found in KV storage") // ErrKVGroupEmpty is returned when a required group name is empty ErrKVGroupEmpty = fmt.Errorf("group name is empty") // ErrKVKeyEmpty is returned when a required key name is empty ErrKVKeyEmpty = fmt.Errorf("key name is empty") // ErrKVKeyAndGroupEmpty is returned when both key and group names are empty ErrKVKeyAndGroupEmpty = fmt.Errorf("both key and group are empty") // ErrKVTransactionFailed is returned when a KV storage transaction operation fails ErrKVTransactionFailed = fmt.Errorf("KV storage transaction failed") ) // KVParams is the parameters for KV storage operations type KVParams struct { Group string Key string Value string Page int PageSize int } // KVOperator provides methods to interact with the key-value storage system for plugins type KVOperator struct { data *Data session *xorm.Session pluginSlugName string cacheTTL time.Duration } // KVStorageOption defines a function type that configures a KVOperator type KVStorageOption func(*KVOperator) // WithCacheTTL is the option to set the cache TTL; the default value is 30 minutes. // If ttl is less than 0, the cache will not be used func WithCacheTTL(ttl time.Duration) KVStorageOption { return func(kv *KVOperator) { kv.cacheTTL = ttl } } // Option is used to set the options for the KV storage func (kv *KVOperator) Option(opts ...KVStorageOption) { for _, opt := range opts { opt(kv) } } func (kv *KVOperator) getSession(ctx context.Context) (*xorm.Session, func()) { session := kv.session cleanup := func() {} if session == nil { session = kv.data.DB.NewSession().Context(ctx) cleanup = func() { if session != nil { _ = session.Close() } } } return session, cleanup } func (kv *KVOperator) getCacheKey(params KVParams) string { return fmt.Sprintf("plugin_kv_storage:%s:group:%s:key:%s", kv.pluginSlugName, params.Group, params.Key) } func (kv *KVOperator) setCache(ctx context.Context, params KVParams) { if kv.cacheTTL < 0 { return } ttl := kv.cacheTTL if ttl > 10 { ttl += time.Duration(float64(ttl) * 0.1 * (1 - rand.Float64())) } cacheKey := kv.getCacheKey(params) if err := kv.data.Cache.SetString(ctx, cacheKey, params.Value, ttl); err != nil { log.Warnf("cache set failed: %v, key: %s", err, cacheKey) } } func (kv *KVOperator) getCache(ctx context.Context, params KVParams) (string, bool, error) { if kv.cacheTTL < 0 { return "", false, nil } cacheKey := kv.getCacheKey(params) return kv.data.Cache.GetString(ctx, cacheKey) } func (kv *KVOperator) cleanCache(ctx context.Context, params KVParams) { if kv.cacheTTL < 0 { return } if err := kv.data.Cache.Del(ctx, kv.getCacheKey(params)); err != nil { log.Warnf("Failed to delete cache for key %s: %v", params.Key, err) } } // Get retrieves a value from KV storage by group and key. // Returns the value as a string or an error if the key is not found. func (kv *KVOperator) Get(ctx context.Context, params KVParams) (string, error) { if params.Key == "" { return "", ErrKVKeyEmpty } if value, exist, err := kv.getCache(ctx, params); err == nil && exist { return value, nil } // query data := entity.PluginKVStorage{} query, cleanup := kv.getSession(ctx) defer cleanup() query.Where(builder.Eq{ "plugin_slug_name": kv.pluginSlugName, "`group`": params.Group, "`key`": params.Key, }) has, err := query.Get(&data) if err != nil { return "", err } if !has { return "", ErrKVKeyNotFound } params.Value = data.Value kv.setCache(ctx, params) return data.Value, nil } // Set stores a value in KV storage with the specified group and key. // Updates the value if it already exists. func (kv *KVOperator) Set(ctx context.Context, params KVParams) error { if params.Key == "" { return ErrKVKeyEmpty } query, cleanup := kv.getSession(ctx) defer cleanup() data := &entity.PluginKVStorage{ PluginSlugName: kv.pluginSlugName, Group: params.Group, Key: params.Key, Value: params.Value, } kv.cleanCache(ctx, params) affected, err := query.Where(builder.Eq{ "plugin_slug_name": kv.pluginSlugName, "`group`": params.Group, "`key`": params.Key, }).Cols("value").Update(data) if err != nil { return err } if affected == 0 { _, err = query.Insert(data) if err != nil { return err } } return nil } // Del removes values from KV storage by group and/or key. // If both group and key are provided, only that specific entry is deleted. // If only group is provided, all entries in that group are deleted. // At least one of group or key must be provided. func (kv *KVOperator) Del(ctx context.Context, params KVParams) error { if params.Key == "" && params.Group == "" { return ErrKVKeyAndGroupEmpty } kv.cleanCache(ctx, params) session, cleanup := kv.getSession(ctx) defer cleanup() session.Where(builder.Eq{ "plugin_slug_name": kv.pluginSlugName, }) if params.Group != "" { session.Where(builder.Eq{"`group`": params.Group}) } if params.Key != "" { session.Where(builder.Eq{"`key`": params.Key}) } _, err := session.Delete(&entity.PluginKVStorage{}) return err } // GetByGroup retrieves all key-value pairs for a specific group with pagination support. // Returns a map of keys to values or an error if the group is empty or not found. func (kv *KVOperator) GetByGroup(ctx context.Context, params KVParams) (map[string]string, error) { if params.Group == "" { return nil, ErrKVGroupEmpty } if params.Page < 1 { params.Page = 1 } if params.PageSize < 1 { params.PageSize = 10 } query, cleanup := kv.getSession(ctx) defer cleanup() var items []entity.PluginKVStorage err := query.Where(builder.Eq{"plugin_slug_name": kv.pluginSlugName, "`group`": params.Group}). Limit(params.PageSize, (params.Page-1)*params.PageSize). OrderBy("id ASC"). Find(&items) if err != nil { return nil, err } result := make(map[string]string, len(items)) for _, item := range items { result[item.Key] = item.Value } return result, nil } // Tx executes a function within a transaction context. If the KVOperator already has a session, // it will use that session. Otherwise, it creates a new transaction session. // The transaction will be committed if the function returns nil, or rolled back if it returns an error. func (kv *KVOperator) Tx(ctx context.Context, fn func(ctx context.Context, kv *KVOperator) error) error { var ( txKv = kv shouldCommit bool ) if kv.session == nil { session := kv.data.DB.NewSession().Context(ctx) if err := session.Begin(); err != nil { _ = session.Close() return fmt.Errorf("%w: begin transaction failed: %v", ErrKVTransactionFailed, err) } defer func() { if !shouldCommit { if rollbackErr := session.Rollback(); rollbackErr != nil { log.Errorf("rollback failed: %v", rollbackErr) } } _ = session.Close() }() txKv = &KVOperator{ session: session, data: kv.data, pluginSlugName: kv.pluginSlugName, } shouldCommit = true } if err := fn(ctx, txKv); err != nil { return fmt.Errorf("%w: %v", ErrKVTransactionFailed, err) } if shouldCommit { if err := txKv.session.Commit(); err != nil { return fmt.Errorf("%w: commit failed: %v", ErrKVTransactionFailed, err) } } return nil } // KVStorage defines the interface for plugins that need data storage capabilities type KVStorage interface { Info() Info SetOperator(operator *KVOperator) } var ( CallKVStorage, registerKVStorage = MakePlugin[KVStorage](true) ) // NewKVOperator creates a new KV storage operator with the specified database engine, cache and plugin name. // It returns a KVOperator instance that can be used to interact with the plugin's storage. func NewKVOperator(db *xorm.Engine, cache cache.Cache, pluginSlugName string) *KVOperator { return &KVOperator{ data: &Data{DB: db, Cache: cache}, pluginSlugName: pluginSlugName, cacheTTL: 30 * time.Minute, } }