323 lines
8.4 KiB
Go
323 lines
8.4 KiB
Go
package index
|
||
|
||
import (
|
||
"bytes"
|
||
"fmt"
|
||
"os"
|
||
"strings"
|
||
"sync"
|
||
|
||
"git.pyer.club/kingecg/godocdb/storage"
|
||
"go.mongodb.org/mongo-driver/bson"
|
||
)
|
||
|
||
// IndexType 索引类型
|
||
type IndexType int
|
||
|
||
const (
|
||
NonUnique IndexType = iota
|
||
Unique
|
||
)
|
||
|
||
// IndexSortOrder 表示索引的排序方式
|
||
type IndexSortOrder string
|
||
|
||
const (
|
||
Ascending IndexSortOrder = "asc"
|
||
Descending IndexSortOrder = "desc"
|
||
)
|
||
|
||
// IndexMetadata 索引元数据
|
||
type IndexMetadata struct {
|
||
Name string
|
||
Type IndexType
|
||
KeyFields []string
|
||
SortOrders []IndexSortOrder // 每个字段对应的排序方式（长度应与KeyFields一致）
|
||
Version int // 索引版本号，用于缓存一致性校验
|
||
}
|
||
|
||
// IndexStore 管理索引存储
|
||
type IndexStore struct {
|
||
storage *storage.LevelDBStorage
|
||
indexes map[string]map[string]string // 缓存索引数据 {indexName: {key: value}}
|
||
mu sync.RWMutex // 保护并发访问
|
||
}
|
||
|
||
// NewIndexStore 创建新的索引存储实例
|
||
func NewIndexStore(path string) (*IndexStore, error) {
|
||
// 确保目录存在
|
||
if strings.HasPrefix(path, "/") {
|
||
path = strings.TrimSuffix(path, "/")
|
||
}
|
||
if !strings.HasSuffix(path, "_index") {
|
||
// add _index suffix to avoid name confilict with document store, do remove this
|
||
path += "_index"
|
||
}
|
||
if err := os.MkdirAll(path, 0755); err != nil {
|
||
return nil, fmt.Errorf("failed to create storage directory: %v", err)
|
||
}
|
||
|
||
storage, err := storage.NewLevelDBStorage(path)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to create LevelDB storage: %v", err)
|
||
}
|
||
|
||
return &IndexStore{
|
||
storage: storage,
|
||
indexes: make(map[string]map[string]string),
|
||
}, nil
|
||
}
|
||
|
||
// CreateIndex 创建索引
|
||
func (is *IndexStore) CreateIndex(indexName string, indexType IndexType, keyFields []string, sortOrders []IndexSortOrder) error {
|
||
is.mu.Lock()
|
||
defer is.mu.Unlock()
|
||
|
||
// 验证keyFields和sortOrders长度一致
|
||
if len(keyFields) != len(sortOrders) {
|
||
return fmt.Errorf("keyFields and sortOrders must have the same length")
|
||
}
|
||
|
||
// 存储索引元数据
|
||
metadata := IndexMetadata{
|
||
Name: indexName,
|
||
Type: indexType,
|
||
KeyFields: keyFields,
|
||
SortOrders: sortOrders,
|
||
Version: 1, // 初始版本号
|
||
}
|
||
|
||
data, err := bson.Marshal(metadata)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to marshal index metadata: %v", err)
|
||
}
|
||
|
||
key := []byte(fmt.Sprintf("indexes:metadata:%s", indexName))
|
||
if err := is.storage.Put(key, data); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 初始化索引结构
|
||
index := make(map[string]string)
|
||
|
||
// 存储初始索引结构
|
||
indexKey := fmt.Sprintf("indexes:data:%s", indexName)
|
||
indexData, err := bson.Marshal(index)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to marshal index data: %v", err)
|
||
}
|
||
|
||
return is.storage.Put([]byte(indexKey), indexData)
|
||
}
|
||
|
||
// UpdateIndex 在指定collection下更新索引
|
||
func (is *IndexStore) UpdateIndex(collection string, indexName string, fieldValue interface{}, documentID string) error {
|
||
is.mu.Lock()
|
||
defer is.mu.Unlock()
|
||
|
||
// 增加空值检查
|
||
if fieldValue == nil {
|
||
return fmt.Errorf("cannot update index with nil field value")
|
||
}
|
||
|
||
// 从存储加载当前索引
|
||
indexKey := []byte(fmt.Sprintf("indexes:data:%s", indexName))
|
||
storedData, err := is.storage.Get(indexKey)
|
||
if err != nil {
|
||
// 如果不存在则创建新索引
|
||
storedData, _ = bson.Marshal(make(map[string]string))
|
||
}
|
||
|
||
// 反序列化索引数据
|
||
var index map[string]string
|
||
if err := bson.Unmarshal(storedData, &index); err != nil {
|
||
return fmt.Errorf("failed to unmarshal index data: %v", err)
|
||
}
|
||
|
||
// 获取索引元数据
|
||
metadata, err := is.GetIndexMetadata(indexName)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get index metadata: %v", err)
|
||
}
|
||
|
||
// 更新索引
|
||
indexKeyField := fmt.Sprintf("%s:%v", collection, fieldValue)
|
||
if _, ok := index[indexKeyField]; ok {
|
||
// 处理重复值的逻辑（需要根据索引类型决定是否覆盖）
|
||
if metadata.Type == Unique {
|
||
return fmt.Errorf("duplicate key error for unique index")
|
||
}
|
||
}
|
||
|
||
index[indexKeyField] = documentID
|
||
|
||
// 序列化并存储更新后的索引
|
||
updatedData, err := bson.Marshal(index)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to marshal updated index: %v", err)
|
||
}
|
||
|
||
return is.storage.Put(indexKey, updatedData)
|
||
}
|
||
|
||
// getOrCreateIndex 获取现有索引或创建新索引
|
||
func (is *IndexStore) getOrCreateIndex(indexName string) (map[string]string, error) {
|
||
// 检查内存缓存
|
||
if index, ok := is.indexes[indexName]; ok {
|
||
return index, nil
|
||
}
|
||
|
||
// 从存储中加载
|
||
indexKey := fmt.Sprintf("indexes:data:%s", indexName)
|
||
storedData, err := is.storage.Get([]byte(indexKey))
|
||
if err != nil {
|
||
// 如果不存在则创建空索引
|
||
return make(map[string]string), nil
|
||
}
|
||
|
||
// 反序列化索引数据
|
||
var loadedIndex map[string]string
|
||
if err := bson.Unmarshal(storedData, &loadedIndex); err != nil {
|
||
return nil, fmt.Errorf("failed to unmarshal index data: %v", err)
|
||
}
|
||
|
||
return loadedIndex, nil
|
||
}
|
||
|
||
// GetIndexedDocuments 查询特定collection下的文档
|
||
func (is *IndexStore) GetIndexedDocuments(collection string, indexName string, fieldValue interface{}) ([]string, error) {
|
||
is.mu.RLock()
|
||
defer is.mu.RUnlock()
|
||
|
||
// 构造前缀
|
||
prefix := fmt.Sprintf("%s:%v", collection, fieldValue)
|
||
|
||
// 直接从存储读取
|
||
indexKey := fmt.Sprintf("indexes:data:%s", indexName)
|
||
storedData, err := is.storage.Get([]byte(indexKey))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("index not found: %v", err)
|
||
}
|
||
|
||
var loadedIndex map[string]string
|
||
if err := bson.Unmarshal(storedData, &loadedIndex); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var results []string
|
||
|
||
// 遍历查找匹配项
|
||
for key, value := range loadedIndex {
|
||
if strings.HasPrefix(key, prefix) {
|
||
results = append(results, value)
|
||
}
|
||
}
|
||
|
||
return results, nil
|
||
}
|
||
|
||
// GetAllIndexMetadata 获取所有索引的元数据
|
||
func (is *IndexStore) GetAllIndexMetadata(collection string) ([]IndexMetadata, error) {
|
||
// 构建前缀键（注意这里移除了多余的冒号）
|
||
prefix := []byte(fmt.Sprintf("indexes:metadata:%s:", collection))
|
||
|
||
// 使用storage的Scan方法替代直接访问db字段
|
||
keys, _, err := is.storage.Scan(prefix)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to scan metadata: %v", err)
|
||
}
|
||
|
||
var metadataList []IndexMetadata
|
||
for _, key := range keys {
|
||
// 验证键是否匹配前缀
|
||
if !bytes.HasPrefix(key, prefix) {
|
||
continue
|
||
}
|
||
|
||
// 读取元数据值
|
||
value, err := is.storage.Get(key)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
|
||
// 解析元数据
|
||
var md IndexMetadata
|
||
if err := bson.Unmarshal(value, &md); err != nil {
|
||
continue
|
||
}
|
||
|
||
metadataList = append(metadataList, md)
|
||
}
|
||
|
||
return metadataList, nil
|
||
}
|
||
|
||
// DeleteIndex 删除索引条目
|
||
func (is *IndexStore) DeleteIndex(collection string, indexName string, fieldValue interface{}, documentID string) error {
|
||
is.mu.Lock()
|
||
defer is.mu.Unlock()
|
||
|
||
// 修复索引键的构造方式，使用统一格式
|
||
indexKey := []byte(fmt.Sprintf("indexes:data:%s:%s:%v", collection, indexName, fieldValue))
|
||
|
||
// 从LevelDB删除
|
||
return is.storage.Delete(indexKey)
|
||
}
|
||
|
||
// DropIndex 删除索引
|
||
func (is *IndexStore) DropIndex(indexName string) error {
|
||
is.mu.Lock()
|
||
defer is.mu.Unlock()
|
||
|
||
// 删除元数据
|
||
metadataKey := []byte(fmt.Sprintf("indexes:metadata:%s", indexName))
|
||
if err := is.storage.Delete(metadataKey); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 删除内存中的缓存
|
||
delete(is.indexes, indexName)
|
||
|
||
// 删除索引数据
|
||
indexKey := []byte(fmt.Sprintf("indexes:data:%s", indexName))
|
||
return is.storage.Delete(indexKey)
|
||
}
|
||
|
||
// GetIndexMetadata 获取索引元数据
|
||
func (is *IndexStore) GetIndexMetadata(indexName string) (IndexMetadata, error) {
|
||
key := []byte(fmt.Sprintf("indexes:metadata:%s", indexName))
|
||
storedData, err := is.storage.Get(key)
|
||
if err != nil {
|
||
return IndexMetadata{}, fmt.Errorf("index not found: %v", err)
|
||
}
|
||
|
||
var metadata IndexMetadata
|
||
if err := bson.Unmarshal(storedData, &metadata); err != nil {
|
||
return IndexMetadata{}, fmt.Errorf("failed to unmarshal index metadata: %v", err)
|
||
}
|
||
|
||
return metadata, nil
|
||
}
|
||
|
||
// FlushIndexToDisk 将内存中的索引更新写入磁盘
|
||
func (is *IndexStore) FlushIndexToDisk(indexName string) error {
|
||
// 因为现在直接操作存储，不需要单独的刷新方法
|
||
// 现在Flush只是简单验证索引是否存在
|
||
is.mu.RLock()
|
||
defer is.mu.RUnlock()
|
||
|
||
if _, ok := is.indexes[indexName]; !ok {
|
||
return fmt.Errorf("index not found in memory: %s", indexName)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// ClearIndexCache 清除内存中的索引缓存
|
||
func (is *IndexStore) ClearIndexCache(indexName string) {
|
||
is.mu.Lock()
|
||
defer is.mu.Unlock()
|
||
|
||
delete(is.indexes, indexName)
|
||
}
|