github.com/kubeshop/testkube@v1.17.23/pkg/repository/result/mongo_numbers.go (about)

     1  package result
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  
     7  	"go.mongodb.org/mongo-driver/bson"
     8  	"go.mongodb.org/mongo-driver/mongo/options"
     9  
    10  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    11  )
    12  
    13  type oldExecutionNumber struct {
    14  	TestName string `json:"testName"`
    15  	Number   int    `json:"number"`
    16  }
    17  
    18  type executionNumber struct {
    19  	Name        string `json:"name"`
    20  	Number      int    `json:"number"`
    21  	IsTestSuite bool   `json:"isTestSuite"`
    22  }
    23  
    24  func (r *MongoRepository) GetNextExecutionNumber(ctx context.Context, name string) (number int32, err error) {
    25  	err = r.convertFromOldToNew()
    26  	if err != nil {
    27  		return 1, err
    28  	}
    29  
    30  	// TODO: modify this when we decide to update the interfaces for OSS and cloud
    31  	isTestSuite := strings.HasPrefix(name, "ts-")
    32  
    33  	execNmbr := executionNumber{Name: name, IsTestSuite: isTestSuite}
    34  	retry := false
    35  	retryAttempts := 0
    36  	maxRetries := 10
    37  
    38  	opts := options.FindOneAndUpdate()
    39  	opts.SetUpsert(true)
    40  	opts.SetReturnDocument(options.After)
    41  
    42  	err = r.SequencesColl.FindOne(ctx, bson.M{"name": name}).Decode(&execNmbr)
    43  	if err != nil {
    44  		var execution testkube.Execution
    45  		number, _ = r.GetLatestTestNumber(ctx, name)
    46  		if number == 0 {
    47  			execNmbr.Number = 1
    48  		} else {
    49  			execNmbr.Number = int(execution.Number) + 1
    50  		}
    51  		_, err = r.SequencesColl.InsertOne(ctx, execNmbr)
    52  	} else {
    53  		err = r.SequencesColl.FindOneAndUpdate(ctx, bson.M{"name": name}, bson.M{"$inc": bson.M{"number": 1}}, opts).Decode(&execNmbr)
    54  	}
    55  
    56  	retry = err != nil
    57  
    58  	for retry {
    59  		retryAttempts++
    60  		err = r.SequencesColl.FindOneAndUpdate(ctx, bson.M{"name": name}, bson.M{"$inc": bson.M{"number": 1}}, opts).Decode(&execNmbr)
    61  		if err == nil || retryAttempts >= maxRetries {
    62  			retry = false
    63  		}
    64  	}
    65  
    66  	return int32(execNmbr.Number), nil
    67  }
    68  
    69  func (r *MongoRepository) DeleteExecutionNumber(ctx context.Context, name string) (err error) {
    70  	err = r.convertFromOldToNew()
    71  	if err != nil {
    72  		return err
    73  	}
    74  	_, err = r.SequencesColl.DeleteOne(ctx, bson.M{"name": name})
    75  	return err
    76  }
    77  
    78  func (r *MongoRepository) DeleteExecutionNumbers(ctx context.Context, names []string) (err error) {
    79  	err = r.convertFromOldToNew()
    80  	if err != nil {
    81  		return err
    82  	}
    83  	_, err = r.SequencesColl.DeleteMany(ctx, bson.M{"name": bson.M{"$in": names}})
    84  	return err
    85  }
    86  
    87  func (r *MongoRepository) DeleteAllExecutionNumbers(ctx context.Context, isTestSuite bool) (err error) {
    88  	err = r.convertFromOldToNew()
    89  	if err != nil {
    90  		return err
    91  	}
    92  	_, err = r.SequencesColl.DeleteMany(ctx, bson.M{"istestsuite": isTestSuite})
    93  	return err
    94  }
    95  
    96  func (r *MongoRepository) convertFromOldToNew() error {
    97  	filter := bson.M{"testname": bson.M{"$exists": true}}
    98  
    99  	cursor, err := r.SequencesColl.Find(context.Background(), filter)
   100  	if err != nil {
   101  		return err
   102  	}
   103  	defer cursor.Close(context.Background())
   104  
   105  	for cursor.Next(context.Background()) {
   106  		var entry oldExecutionNumber
   107  		err := cursor.Decode(&entry)
   108  		if err != nil {
   109  			return err
   110  		}
   111  
   112  		isTestSuite := strings.HasPrefix(entry.TestName, "ts-")
   113  
   114  		newEntry := executionNumber{
   115  			Name:        entry.TestName,
   116  			Number:      entry.Number,
   117  			IsTestSuite: isTestSuite,
   118  		}
   119  
   120  		_, err = r.SequencesColl.InsertOne(context.Background(), newEntry)
   121  		if err != nil {
   122  			return err
   123  		}
   124  	}
   125  	_, err = r.SequencesColl.DeleteMany(context.Background(), filter)
   126  	return err
   127  }