keepalive/adapter/mongodb.go

package adapter

import (
	"context"

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

func init() {
	Registry["mongodb"] = func(cfg Config) (Adapter, error) {
		uri, err := cfg.Required("uri")
		if err != nil {
			return nil, err
		}
		dbName, err := cfg.Required("database")
		if err != nil {
			return nil, err
		}
		collName, err := cfg.Required("collection")
		if err != nil {
			return nil, err
		}
		return &mongoAdapter{
			uri:      uri,
			dbName:   dbName,
			collName: collName,
			docID:    cfg.Optional("counter_key", "counter"),
		}, nil
	}
	Registry["mongo"] = Registry["mongodb"]
}

type mongoAdapter struct {
	client   *mongo.Client
	coll     *mongo.Collection
	uri      string
	dbName   string
	collName string
	docID    string
}

func (a *mongoAdapter) Connect(ctx context.Context) error {
	client, err := mongo.Connect(options.Client().ApplyURI(a.uri))
	if err != nil {
		return err
	}
	a.client = client
	a.coll = client.Database(a.dbName).Collection(a.collName)
	if err := client.Ping(ctx, nil); err != nil {
		client.Disconnect(ctx)
		return err
	}
	if err := a.ensureInitialized(ctx); err != nil {
		client.Disconnect(ctx)
		return err
	}
	return nil
}

func (a *mongoAdapter) ensureInitialized(ctx context.Context) error {
	update := bson.M{"$setOnInsert": bson.M{"count": int64(0)}}
	_, err := a.coll.UpdateOne(ctx, bson.M{"_id": a.docID}, update, options.UpdateOne().SetUpsert(true))
	return err
}

func (a *mongoAdapter) Increment(ctx context.Context) (int64, error) {
	update := bson.M{"$inc": bson.M{"count": 1}}
	opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)
	var out struct {
		Count int64 `bson:"count"`
	}
	err := a.coll.FindOneAndUpdate(ctx, bson.M{"_id": a.docID}, update, opts).Decode(&out)
	if err != nil {
		return 0, err
	}
	return out.Count, nil
}

func (a *mongoAdapter) Close(ctx context.Context) error {
	if a.client == nil {
		return nil
	}
	return a.client.Disconnect(ctx)
}