github.com/mdaxf/iac@v0.0.0-20240519030858-58a061660378/health/checks/mongodb.go (about)

     1  package checks
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"go.mongodb.org/mongo-driver/mongo"
     9  	"go.mongodb.org/mongo-driver/mongo/options"
    10  	"go.mongodb.org/mongo-driver/mongo/readpref"
    11  )
    12  
    13  const (
    14  	defaultTimeoutConnect    = 5 * time.Second
    15  	defaultTimeoutDisconnect = 5 * time.Second
    16  	defaultTimeoutPing       = 5 * time.Second
    17  )
    18  
    19  // Config is the MongoDB checker configuration settings container.
    20  type MongoDbCheck struct {
    21  	Ctx context.Context
    22  
    23  	MongoDBClient *mongo.Client
    24  	// DSN is the MongoDB instance connection DSN. Required.
    25  	ConnectionString string
    26  
    27  	// TimeoutConnect defines timeout for establishing mongo connection, if not set - default value is used
    28  	TimeoutConnect time.Duration
    29  	// TimeoutDisconnect defines timeout for closing connection, if not set - default value is used
    30  	TimeoutDisconnect time.Duration
    31  	// TimeoutDisconnect defines timeout for making ping request, if not set - default value is used
    32  	TimeoutPing time.Duration
    33  	// return the check error
    34  	Error error
    35  }
    36  
    37  func CheckMongoDBStatus(ctx context.Context, connectionString string, timeoutConnect time.Duration, timeoutDisconnect time.Duration, timeoutPing time.Duration) error {
    38  	check := NewMongoDbCheck(ctx, connectionString, timeoutConnect, timeoutDisconnect, timeoutPing)
    39  	return check.CheckStatus()
    40  }
    41  
    42  func NewMongoDbCheck(ctx context.Context, connectionString string, timeoutConnect time.Duration, timeoutDisconnect time.Duration, timeoutPing time.Duration) MongoDbCheck {
    43  	if timeoutConnect == 0 {
    44  		timeoutConnect = defaultTimeoutConnect
    45  	}
    46  
    47  	if timeoutDisconnect == 0 {
    48  		timeoutDisconnect = defaultTimeoutDisconnect
    49  	}
    50  
    51  	if timeoutPing == 0 {
    52  		timeoutPing = defaultTimeoutPing
    53  	}
    54  
    55  	return MongoDbCheck{
    56  		Ctx:               ctx,
    57  		ConnectionString:  connectionString,
    58  		Error:             nil,
    59  		TimeoutConnect:    timeoutConnect,
    60  		TimeoutDisconnect: timeoutDisconnect,
    61  		TimeoutPing:       timeoutPing,
    62  	}
    63  }
    64  
    65  func (check MongoDbCheck) CheckStatus() error {
    66  
    67  	defer func() {
    68  		if r := recover(); r != nil {
    69  			return
    70  		}
    71  	}()
    72  
    73  	var checkErr error
    74  	checkErr = nil
    75  	ctx := check.Ctx
    76  	client, err := mongo.NewClient(options.Client().ApplyURI(check.ConnectionString))
    77  	if err != nil {
    78  		checkErr = fmt.Errorf("mongoDB health check failed on client creation: %w", err)
    79  		check.Error = checkErr
    80  		return checkErr
    81  	}
    82  
    83  	ctxConn, cancelConn := context.WithTimeout(ctx, check.TimeoutConnect)
    84  	defer cancelConn()
    85  
    86  	err = client.Connect(ctxConn)
    87  	if err != nil {
    88  		checkErr = fmt.Errorf("mongoDB health check failed on connect: %w", err)
    89  		check.Error = checkErr
    90  		return checkErr
    91  	}
    92  
    93  	defer func() {
    94  		ctxDisc, cancelDisc := context.WithTimeout(ctx, check.TimeoutDisconnect)
    95  		defer cancelDisc()
    96  
    97  		// override checkErr only if there were no other errors
    98  		if err := client.Disconnect(ctxDisc); err != nil && checkErr == nil {
    99  			checkErr = fmt.Errorf("mongoDB health check failed on closing connection: %w", err)
   100  			check.Error = checkErr
   101  			return
   102  		}
   103  	}()
   104  
   105  	ctxPing, cancelPing := context.WithTimeout(ctx, check.TimeoutPing)
   106  	defer cancelPing()
   107  
   108  	err = client.Ping(ctxPing, readpref.Primary())
   109  	if err != nil {
   110  		checkErr = fmt.Errorf("mongoDB health check failed on ping: %w", err)
   111  		check.Error = checkErr
   112  		return checkErr
   113  	}
   114  
   115  	return nil
   116  }