wehub-resource-sync
337 lines
9.0 KiB
Go
337 lines
9.0 KiB
Go
/*
|
|
* 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,
|
|
}
|
|
}
|