gotidb/pkg/engine/bolt/bolt_query.go

package bolt

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"time"

	"git.pyer.club/kingecg/gotidb/pkg/engine"
	bolt "go.etcd.io/bbolt"
)

// Query 实现 Engine 接口
func (b *BoltEngine) Query(ctx context.Context, query engine.Query) (engine.QueryResult, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if !b.opened || b.closed {
		return nil, fmt.Errorf("bolt engine not open")
	}

	startTime := time.Now()
	var result engine.QueryResult

	err := b.db.View(func(tx *bolt.Tx) error {
		// 遍历所有序列
		indexBucket := tx.Bucket([]byte(indexBucketName))
		if indexBucket == nil {
			return fmt.Errorf("index bucket not found")
		}

		return indexBucket.ForEach(func(k, v []byte) error {
			seriesID := string(k)
			bucket := tx.Bucket([]byte(seriesBucketPrefix + seriesID))
			if bucket == nil {
				return nil
			}

			// 根据查询类型执行不同的查询
			var seriesResult *engine.SeriesResult
			var err error

			switch query.Type {
			case engine.QueryTypeLatest:
				seriesResult, err = b.queryLatest(bucket, seriesID, query)
			case engine.QueryTypeRaw:
				seriesResult, err = b.queryRaw(bucket, seriesID, query)
			case engine.QueryTypeAggregate:
				seriesResult, err = b.queryAggregate(bucket, seriesID, query)
			case engine.QueryTypeValueDuration:
				seriesResult, err = b.queryValueDuration(bucket, seriesID, query)
			default:
				return fmt.Errorf("unsupported query type: %s", query.Type)
			}

			if err != nil {
				return err
			}

			if seriesResult != nil {
				result = append(result, *seriesResult)
			}

			return nil
		})
	})

	if err != nil {
		b.stats.QueryErrors++
		return nil, fmt.Errorf("failed to execute query: %v", err)
	}

	// 更新统计信息
	b.stats.QueryLatency = time.Since(startTime)

	return result, nil
}

// queryLatest 执行最新值查询
func (b *BoltEngine) queryLatest(bucket *bolt.Bucket, seriesID string, query engine.Query) (*engine.SeriesResult, error) {
	cursor := bucket.Cursor()
	k, v := cursor.Last()
	if k == nil {
		return nil, nil
	}

	var point engine.DataPoint
	if err := json.Unmarshal(v, &point); err != nil {
		return nil, fmt.Errorf("failed to unmarshal data point: %v", err)
	}

	// 检查标签是否匹配
	if !matchTags(point.Labels, query.Tags) {
		return nil, nil
	}

	return &engine.SeriesResult{
		SeriesID: seriesID,
		Points:   []engine.DataPoint{point},
	}, nil
}

// queryRaw 执行原始数据查询
func (b *BoltEngine) queryRaw(bucket *bolt.Bucket, seriesID string, query engine.Query) (*engine.SeriesResult, error) {
	var points []engine.DataPoint

	cursor := bucket.Cursor()
	for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
		timestamp := int64(binary.BigEndian.Uint64(k))
		if timestamp < query.StartTime {
			continue
		}
		if timestamp > query.EndTime {
			break
		}

		var point engine.DataPoint
		if err := json.Unmarshal(v, &point); err != nil {
			return nil, fmt.Errorf("failed to unmarshal data point: %v", err)
		}

		if len(points) == 0 {
			// 检查标签是否匹配
			if !matchTags(point.Labels, query.Tags) {
				return nil, nil
			}
		}

		points = append(points, point)

		// 应用限制
		if query.Limit > 0 && len(points) >= query.Limit {
			break
		}
	}

	if len(points) == 0 {
		return nil, nil
	}

	return &engine.SeriesResult{
		SeriesID: seriesID,
		Points:   points,
	}, nil
}

// queryAggregate 执行聚合查询
func (b *BoltEngine) queryAggregate(bucket *bolt.Bucket, seriesID string, query engine.Query) (*engine.SeriesResult, error) {
	var points []engine.DataPoint
	var firstPoint engine.DataPoint

	cursor := bucket.Cursor()
	for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
		timestamp := int64(binary.BigEndian.Uint64(k))
		if timestamp < query.StartTime {
			continue
		}
		if timestamp > query.EndTime {
			break
		}

		var point engine.DataPoint
		if err := json.Unmarshal(v, &point); err != nil {
			return nil, fmt.Errorf("failed to unmarshal data point: %v", err)
		}

		if len(points) == 0 {
			firstPoint = point
			// 检查标签是否匹配
			if !matchTags(point.Labels, query.Tags) {
				return nil, nil
			}
		}

		points = append(points, point)
	}

	if len(points) == 0 {
		return nil, nil
	}

	// 计算聚合值
	aggregateValue := calculateAggregate(points, query.AggregateType)

	// 创建聚合结果点
	aggregatePoint := engine.DataPoint{
		Timestamp: query.EndTime,
		Value:     aggregateValue,
		Labels:    firstPoint.Labels,
	}

	return &engine.SeriesResult{
		SeriesID: seriesID,
		Points:   []engine.DataPoint{aggregatePoint},
	}, nil
}

// queryValueDuration 执行值持续时间查询
func (b *BoltEngine) queryValueDuration(bucket *bolt.Bucket, seriesID string, query engine.Query) (*engine.SeriesResult, error) {
	// 简化实现，实际应该计算每个值的持续时间
	return b.queryRaw(bucket, seriesID, query)
}