sendico/api/pkg/db/internal/mongo/tseriesimp/tseries.go

package tseriesimp

import (
	"context"
	"errors"
	"time"

	"github.com/tech/sendico/pkg/db/repository"
	"github.com/tech/sendico/pkg/db/repository/builder"
	rdecoder "github.com/tech/sendico/pkg/db/repository/decoder"
	tsoptions "github.com/tech/sendico/pkg/db/tseries/options"
	tspoint "github.com/tech/sendico/pkg/db/tseries/point"
	"github.com/tech/sendico/pkg/merrors"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type TimeSeries struct {
	options    tsoptions.Options
	collection *mongo.Collection
}

func NewMongoTimeSeriesCollection(ctx context.Context, db *mongo.Database, tsOpts *tsoptions.Options) (*TimeSeries, error) {
	if tsOpts == nil {
		return nil, merrors.InvalidArgument("nil time-series options provided", "options")
	}
	// Configure time-series options
	granularity := tsOpts.Granularity.String()
	ts := &options.TimeSeriesOptions{
		TimeField:   tsOpts.TimeField,
		Granularity: &granularity,
	}
	if tsOpts.MetaField != "" {
		ts.MetaField = &tsOpts.MetaField
	}

	// Collection options
	collOpts := options.CreateCollection().SetTimeSeriesOptions(ts)

	// Set TTL if requested
	if tsOpts.ExpireAfter > 0 {
		secs := int64(tsOpts.ExpireAfter / time.Second)
		collOpts.SetExpireAfterSeconds(secs)
	}

	if err := db.CreateCollection(ctx, tsOpts.Collection, collOpts); err != nil {
		if cmdErr, ok := err.(mongo.CommandError); !ok || cmdErr.Code != 48 {
			return nil, err
		}
	}

	return &TimeSeries{collection: db.Collection(tsOpts.Collection), options: *tsOpts}, nil
}

func (ts *TimeSeries) Aggregate(ctx context.Context, pipeline builder.Pipeline, decoder rdecoder.DecodingFunc) error {
	queryFunc := func(ctx context.Context, collection *mongo.Collection) (*mongo.Cursor, error) {
		return collection.Aggregate(ctx, pipeline.Build())
	}
	return ts.executeQuery(ctx, decoder, queryFunc)
}

func (ts *TimeSeries) Insert(ctx context.Context, timePoint tspoint.TimePoint) error {
	_, err := ts.collection.InsertOne(ctx, timePoint)
	return err
}

func (ts *TimeSeries) InsertMany(ctx context.Context, timePoints []tspoint.TimePoint) error {
	docs := make([]any, len(timePoints))
	for i, p := range timePoints {
		docs[i] = p
	}

	// ignore the result if you like, or capture it
	_, err := ts.collection.InsertMany(ctx, docs)
	return err
}

type QueryFunc func(ctx context.Context, collection *mongo.Collection) (*mongo.Cursor, error)

func (ts *TimeSeries) executeQuery(ctx context.Context, decoder rdecoder.DecodingFunc, queryFunc QueryFunc) error {
	cursor, err := queryFunc(ctx, ts.collection)
	if errors.Is(err, mongo.ErrNoDocuments) {
		return merrors.NoData("no_items_in_array")
	}
	if err != nil {
		return err
	}
	defer cursor.Close(ctx)

	for cursor.Next(ctx) {
		if err := cursor.Err(); err != nil {
			return err
		}
		if err = decoder(cursor); err != nil {
			return err
		}
	}

	return nil
}

func (ts *TimeSeries) Query(ctx context.Context, decoder rdecoder.DecodingFunc, query builder.Query, from, to *time.Time) error {
	timeLimitedQuery := query
	if from != nil {
		timeLimitedQuery = timeLimitedQuery.And(repository.Query().Comparison(repository.Field(ts.options.TimeField), builder.Gte, *from))
	}
	if to != nil {
		timeLimitedQuery = timeLimitedQuery.And(repository.Query().Comparison(repository.Field(ts.options.TimeField), builder.Lte, *to))
	}
	queryFunc := func(ctx context.Context, collection *mongo.Collection) (*mongo.Cursor, error) {
		return collection.Find(ctx, timeLimitedQuery.BuildQuery(), timeLimitedQuery.BuildOptions())
	}
	return ts.executeQuery(ctx, decoder, queryFunc)
}

func (ts *TimeSeries) Name() string {
	return ts.collection.Name()
}