zgene
6bed393c12
Backend Tests / backend-unit-test (push) Has been cancelled
Backend Tests / benchmark-test (push) Has been cancelled
CI@main / Node.js v22 (ubuntu-latest) (push) Has been cancelled
Thrift Syntax Validation / validate-thrift (push) Has been cancelled
License Check / License Check (push) Has been cancelled
119 lines
2.8 KiB
Go
119 lines
2.8 KiB
Go
/*
|
|
* Copyright 2025 coze-dev Authors
|
|
*
|
|
* Licensed 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 kvstore
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var ErrKeyNotFound = errors.New("key not found")
|
|
|
|
/*
|
|
CREATE TABLE IF NOT EXISTS `kv_entries` (
|
|
`id` BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
|
|
`namespace` VARCHAR(255) NOT NULL,
|
|
`key_data` VARCHAR(255) NOT NULL,
|
|
`value_data` LONGBLOB NOT NULL,
|
|
PRIMARY KEY (`id`),
|
|
UNIQUE KEY `uniq_namespace_key` (`namespace`, `key_data`)
|
|
) ENGINE=InnoDB CHARSET utf8mb4 COLLATE utf8mb4_general_ci COMMENT 'kv data';
|
|
*/
|
|
|
|
type KVStore[T any] struct {
|
|
repo *gorm.DB
|
|
}
|
|
|
|
var defaultDB *gorm.DB
|
|
|
|
func SetDefault(db *gorm.DB) {
|
|
defaultDB = db
|
|
}
|
|
|
|
func New[T any](db *gorm.DB) *KVStore[T] {
|
|
return &KVStore[T]{
|
|
repo: db,
|
|
}
|
|
}
|
|
|
|
func (g *KVStore[T]) db(ctx context.Context) *gorm.DB {
|
|
if g.repo == nil {
|
|
return defaultDB.WithContext(ctx)
|
|
}
|
|
|
|
return g.repo.WithContext(ctx)
|
|
}
|
|
|
|
func (g *KVStore[T]) Save(ctx context.Context, namespace, k string, v *T) error {
|
|
if v == nil {
|
|
return fmt.Errorf("cannot save nil value for key: %s", k)
|
|
}
|
|
|
|
data, err := json.Marshal(v)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal failed for key %s for type %T: %w", k, *v, err)
|
|
}
|
|
|
|
res := g.db(ctx).Exec(
|
|
"INSERT INTO `kv_entries` (`namespace`, `key_data`, `value_data`) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE `value_data` = ?",
|
|
namespace, k, data, data,
|
|
)
|
|
|
|
if res.Error != nil {
|
|
return fmt.Errorf("failed to save key %s: %w", k, res.Error)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (g *KVStore[T]) Get(ctx context.Context, namespace, k string) (*T, error) {
|
|
var obj T
|
|
|
|
row := g.db(ctx).Raw(
|
|
"SELECT `value_data` FROM `kv_entries` WHERE `namespace` = ? AND `key_data` = ? LIMIT 1",
|
|
namespace, k,
|
|
).Row()
|
|
|
|
var value []byte
|
|
if err := row.Scan(&value); err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil, ErrKeyNotFound
|
|
}
|
|
return nil, fmt.Errorf("failed to get key %s: %w", k, err)
|
|
}
|
|
|
|
if err := json.Unmarshal(value, &obj); err != nil {
|
|
return nil, fmt.Errorf("failed to unmarshal json for key %s: %w", k, err)
|
|
}
|
|
|
|
return &obj, nil
|
|
}
|
|
|
|
func (g *KVStore[T]) Delete(ctx context.Context, namespace, k string) error {
|
|
res := g.db(ctx).Exec(
|
|
"DELETE FROM `kv_entries` WHERE `namespace` = ? AND `key_data` = ?",
|
|
namespace, k,
|
|
)
|
|
|
|
return res.Error
|
|
}
|