sxchnge/backend/shared/database/internal/repository/repository.go

package repositoryimp

import (
	"context"
	"errors"
	"fmt"

	"git.sendico.io/tech/sxchnge/shared/database/record"
	"git.sendico.io/tech/sxchnge/shared/database/repository/builder"
	rd "git.sendico.io/tech/sxchnge/shared/database/repository/decoder"
	"git.sendico.io/tech/sxchnge/shared/database/source"
	"git.sendico.io/tech/sxchnge/shared/database/transaction"
	"git.sendico.io/tech/sxchnge/shared/errors"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type MongoRepository struct {
	collectionName string
	source         source.Source
}

func idFilter(id bson.ObjectID) bson.D {
	return bson.D{
		{Key: storable.IDField, Value: id},
	}
}

func NewMongoRepository(db *mongo.Database, collection string) *MongoRepository {
	return NewRepository(source.Fixed(db), collection)
}

func NewRepository(database source.Source, collection string) *MongoRepository {
	return &MongoRepository{collectionName: collection, source: database}
}

func (r *MongoRepository) acquire(ctx context.Context) (*mongo.Collection, context.Context, func(), error) {
	if r.source == nil {
		return nil, ctx, nil, merrors.FailedPrecondition("database source is not configured")
	}

	database, scoped, release, err := r.source.Acquire(ctx)
	if err != nil {
		return nil, ctx, nil, err
	}
	return database.Collection(r.collectionName), scoped, release, nil
}

func (r *MongoRepository) Collection() string {
	return r.collectionName
}

func prepareForWrite(obj storable.Storable) bson.ObjectID {
	id := obj.GetID()
	if id == nil || id.IsZero() {
		obj.SetID(bson.NewObjectID())
		id = obj.GetID()
	}
	obj.Update()
	return *id
}

func (r *MongoRepository) Insert(ctx context.Context, obj storable.Storable, getFilter builder.Query) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	prepareForWrite(obj)
	_, err = collection.InsertOne(ctx, obj)
	if mongo.IsDuplicateKeyError(err) {
		if getFilter != nil {
			if err = r.FindOneByFilter(ctx, getFilter, obj); err != nil {
				return err
			}
		}
		return merrors.DataConflict("duplicate_key")
	}
	return err
}

func (r *MongoRepository) InsertMany(ctx context.Context, objects []storable.Storable) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	if len(objects) == 0 {
		return nil
	}

	docs := make([]interface{}, len(objects))
	for i, obj := range objects {
		prepareForWrite(obj)
		docs[i] = obj
	}

	_, err = collection.InsertMany(ctx, docs)
	return err
}

func (r *MongoRepository) findOneByFilterImp(ctx context.Context, filter bson.D, errMessage string, result storable.Storable) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	err = collection.FindOne(ctx, filter).Decode(result)
	if errors.Is(err, mongo.ErrNoDocuments) {
		return merrors.NoData(errMessage)
	}
	return err
}

func (r *MongoRepository) Get(ctx context.Context, id bson.ObjectID, result storable.Storable) error {
	if id.IsZero() {
		return merrors.InvalidArgument("zero id provided while fetching "+result.Collection(), "id")
	}
	return r.findOneByFilterImp(ctx, idFilter(id), fmt.Sprintf("%s with ID = %s not found", result.Collection(), id.Hex()), result)
}

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

func (r *MongoRepository) executeQuery(ctx context.Context, queryFunc QueryFunc, decoder rd.DecodingFunc) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	cursor, err := queryFunc(ctx, collection)
	if errors.Is(err, mongo.ErrNoDocuments) {
		return merrors.NoData("no_items_in_array")
	}
	if err != nil {
		return err
	}
	defer func() {
		_ = cursor.Close(ctx)
	}()

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

	return cursor.Err()
}

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

func (r *MongoRepository) FindManyByFilter(ctx context.Context, query builder.Query, decoder rd.DecodingFunc) error {
	queryFunc := func(ctx context.Context, collection *mongo.Collection) (*mongo.Cursor, error) {
		return collection.Find(ctx, query.BuildQuery(), query.BuildOptions())
	}
	return r.executeQuery(ctx, queryFunc, decoder)
}

func (r *MongoRepository) FindAllByFilter(ctx context.Context, query builder.Query, result any) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	if query == nil {
		cursor, err := collection.Find(ctx, bson.D{})
		if err != nil {
			return err
		}
		defer func() {
			_ = cursor.Close(ctx)
		}()
		return cursor.All(ctx, result)
	}
	cursor, err := collection.Find(ctx, query.BuildQuery(), query.BuildOptions())
	if err != nil {
		return err
	}
	defer func() {
		_ = cursor.Close(ctx)
	}()
	return cursor.All(ctx, result)
}

func (r *MongoRepository) FindOneByFilter(ctx context.Context, query builder.Query, result storable.Storable) error {
	return r.findOneByFilterImp(ctx, query.BuildQuery(), result.Collection()+" not found by filter", result)
}

func (r *MongoRepository) Update(ctx context.Context, obj storable.Storable) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	obj.Update()
	return collection.FindOneAndReplace(ctx, idFilter(*obj.GetID()), obj).Err()
}

func (r *MongoRepository) ReplaceOneByFilter(ctx context.Context, query builder.Query, obj storable.Storable) (bool, error) {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return false, err
	}
	defer release()

	if query == nil {
		return false, merrors.InvalidArgument("filter is required", "filter")
	}
	if obj == nil {
		return false, merrors.InvalidArgument("object is required", "obj")
	}
	result, err := collection.ReplaceOne(ctx, query.BuildQuery(), obj)
	if mongo.IsDuplicateKeyError(err) {
		return false, merrors.DataConflict("duplicate_key")
	}
	if err != nil {
		return false, err
	}
	return result.MatchedCount > 0, nil
}

func (r *MongoRepository) Upsert(ctx context.Context, obj storable.Storable) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	id := prepareForWrite(obj)
	_, err = collection.ReplaceOne(ctx, idFilter(id), obj, options.Replace().SetUpsert(true))
	return err
}

func (r *MongoRepository) Patch(ctx context.Context, id bson.ObjectID, patch builder.Patch) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	if id.IsZero() {
		return merrors.InvalidArgument("zero id provided while patching", "id")
	}
	_, err = collection.UpdateByID(ctx, id, patch.Build())
	return err
}

func (r *MongoRepository) PatchMany(ctx context.Context, query builder.Query, patch builder.Patch) (int, error) {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return 0, err
	}
	defer release()

	result, err := collection.UpdateMany(ctx, query.BuildQuery(), patch.Build())
	if err != nil {
		return 0, err
	}
	return int(result.ModifiedCount), nil
}

func (r *MongoRepository) ListIDs(ctx context.Context, query builder.Query) ([]bson.ObjectID, error) {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer release()

	filter := query.BuildQuery()
	findOptions := query.BuildOptions().SetProjection(bson.M{storable.IDField: 1})

	cursor, err := collection.Find(ctx, filter, findOptions)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = cursor.Close(ctx)
	}()

	var ids []bson.ObjectID
	for cursor.Next(ctx) {
		var doc struct {
			ID bson.ObjectID `bson:"_id"`
		}
		if err := cursor.Decode(&doc); err != nil {
			return nil, err
		}
		ids = append(ids, doc.ID)
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}

	return ids, nil
}

func (r *MongoRepository) Delete(ctx context.Context, id bson.ObjectID) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	_, err = collection.DeleteOne(ctx, idFilter(id))
	return err
}

func (r *MongoRepository) DeleteMany(ctx context.Context, query builder.Query) error {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return err
	}
	defer release()

	_, err = collection.DeleteMany(ctx, query.BuildQuery())
	return err
}

func (r *MongoRepository) Name() string {
	return r.collectionName
}

func (r *MongoRepository) WithTransaction(ctx context.Context, cb transaction.Callback) (any, error) {
	collection, ctx, release, err := r.acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer release()

	if cb == nil {
		return nil, merrors.InvalidArgument("transaction callback is required", "callback")
	}
	if transaction.ActiveMongoTransaction(ctx) {
		return cb(ctx)
	}
	session, err := collection.Database().Client().StartSession()
	if err != nil {
		return nil, err
	}
	defer session.EndSession(ctx)
	return session.WithTransaction(ctx, cb)
}