github.com/kubeshop/testkube@v1.17.23/pkg/repository/config/mongo.go (about)

     1  package config
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"go.mongodb.org/mongo-driver/bson"
     8  	"go.mongodb.org/mongo-driver/mongo"
     9  	"go.mongodb.org/mongo-driver/mongo/options"
    10  
    11  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    12  	"github.com/kubeshop/testkube/pkg/telemetry"
    13  )
    14  
    15  const CollectionName = "config"
    16  const Id = "api"
    17  
    18  func NewMongoRepository(db *mongo.Database, opts ...Opt) *MongoRepository {
    19  	r := &MongoRepository{
    20  		Coll: db.Collection(CollectionName),
    21  	}
    22  
    23  	for _, opt := range opts {
    24  		opt(r)
    25  	}
    26  
    27  	return r
    28  }
    29  
    30  type Opt func(*MongoRepository)
    31  
    32  func WithMongoRepositoryCollection(collection *mongo.Collection) Opt {
    33  	return func(r *MongoRepository) {
    34  		r.Coll = collection
    35  	}
    36  }
    37  
    38  type MongoRepository struct {
    39  	Coll *mongo.Collection
    40  }
    41  
    42  func (r *MongoRepository) GetUniqueClusterId(ctx context.Context) (clusterId string, err error) {
    43  	config := testkube.Config{}
    44  	_ = r.Coll.FindOne(ctx, bson.M{"id": Id}).Decode(&config)
    45  
    46  	// generate new cluster id and save if there is not already
    47  	if config.ClusterId == "" {
    48  		config.ClusterId = fmt.Sprintf("cluster%s", telemetry.GetMachineID())
    49  		_, err := r.Upsert(ctx, config)
    50  		return config.ClusterId, err
    51  	}
    52  
    53  	return config.ClusterId, nil
    54  }
    55  
    56  func (r *MongoRepository) GetTelemetryEnabled(ctx context.Context) (ok bool, err error) {
    57  	config := testkube.Config{}
    58  	err = r.Coll.FindOne(ctx, bson.M{"id": Id}).Decode(&config)
    59  	return config.EnableTelemetry, err
    60  }
    61  
    62  func (r *MongoRepository) Get(ctx context.Context) (result testkube.Config, err error) {
    63  	err = r.Coll.FindOne(ctx, bson.M{"id": Id}).Decode(&result)
    64  	return
    65  }
    66  
    67  func (r *MongoRepository) Upsert(ctx context.Context, result testkube.Config) (updated testkube.Config, err error) {
    68  	upsert := true
    69  	result.Id = Id
    70  	_, err = r.Coll.ReplaceOne(ctx, bson.M{"id": Id}, result, &options.ReplaceOptions{Upsert: &upsert})
    71  	return result, err
    72  }