github.heygears.com/openimsdk/tools@v0.0.49/db/mongoutil/mongo.go (about)

     1  // Copyright © 2024 OpenIM open source community. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package mongoutil
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"github.com/openimsdk/tools/db/tx"
    22  	"github.com/openimsdk/tools/errs"
    23  	"github.com/openimsdk/tools/mw/specialerror"
    24  	"go.mongodb.org/mongo-driver/mongo"
    25  	"go.mongodb.org/mongo-driver/mongo/options"
    26  )
    27  
    28  func init() {
    29  	if err := specialerror.AddReplace(mongo.ErrNoDocuments, errs.ErrRecordNotFound); err != nil {
    30  		panic(err)
    31  	}
    32  }
    33  
    34  // Config represents the MongoDB configuration.
    35  type Config struct {
    36  	Uri         string
    37  	Address     []string
    38  	Database    string
    39  	Username    string
    40  	Password    string
    41  	MaxPoolSize int
    42  	MaxRetry    int
    43  }
    44  
    45  type Client struct {
    46  	tx tx.Tx
    47  	db *mongo.Database
    48  }
    49  
    50  func (c *Client) GetDB() *mongo.Database {
    51  	return c.db
    52  }
    53  
    54  func (c *Client) GetTx() tx.Tx {
    55  	return c.tx
    56  }
    57  
    58  // NewMongoDB initializes a new MongoDB connection.
    59  func NewMongoDB(ctx context.Context, config *Config) (*Client, error) {
    60  	if err := config.ValidateAndSetDefaults(); err != nil {
    61  		return nil, err
    62  	}
    63  	opts := options.Client().ApplyURI(config.Uri).SetMaxPoolSize(uint64(config.MaxPoolSize))
    64  	var (
    65  		cli *mongo.Client
    66  		err error
    67  	)
    68  	for i := 0; i < config.MaxRetry; i++ {
    69  		cli, err = connectMongo(ctx, opts)
    70  		if err != nil && shouldRetry(ctx, err) {
    71  			time.Sleep(time.Second / 2)
    72  			continue
    73  		}
    74  		break
    75  	}
    76  	if err != nil {
    77  		return nil, errs.WrapMsg(err, "failed to connect to MongoDB", "URI", config.Uri)
    78  	}
    79  	mtx, err := NewMongoTx(ctx, cli)
    80  	if err != nil {
    81  		return nil, err
    82  	}
    83  	return &Client{
    84  		tx: mtx,
    85  		db: cli.Database(config.Database),
    86  	}, nil
    87  }
    88  
    89  func connectMongo(ctx context.Context, opts *options.ClientOptions) (*mongo.Client, error) {
    90  	cli, err := mongo.Connect(ctx, opts)
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  	if err := cli.Ping(ctx, nil); err != nil {
    95  		return nil, err
    96  	}
    97  	return cli, nil
    98  }