withchao
2 years ago
parent
a8c0250944
commit
7ee9e4cdde
@ -1,109 +0,0 @@
|
||||
// Copyright © 2023 OpenIM. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/openimsdk/open-im-server/v3/pkg/common/config"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
"github.com/openimsdk/tools/log"
|
||||
"github.com/openimsdk/tools/mw/specialerror"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var (
|
||||
// Singleton pattern.
|
||||
redisClient redis.UniversalClient
|
||||
)
|
||||
|
||||
const (
|
||||
maxRetry = 10 // number of retries
|
||||
)
|
||||
|
||||
// NewRedis Initialize redis connection.
|
||||
func NewRedis(ctx context.Context, redisConf *config.Redis) (redis.UniversalClient, error) {
|
||||
if redisClient != nil {
|
||||
return redisClient, nil
|
||||
}
|
||||
|
||||
// Read configuration from environment variables
|
||||
overrideConfigFromEnv(redisConf)
|
||||
|
||||
if len(redisConf.Address) == 0 {
|
||||
return nil, errs.Wrap(errors.New("redis address is empty"))
|
||||
}
|
||||
specialerror.AddReplace(redis.Nil, errs.ErrRecordNotFound)
|
||||
var rdb redis.UniversalClient
|
||||
if len(redisConf.Address) > 1 || redisConf.ClusterMode {
|
||||
rdb = redis.NewClusterClient(&redis.ClusterOptions{
|
||||
Addrs: redisConf.Address,
|
||||
Username: redisConf.Username,
|
||||
Password: redisConf.Password, // no password set
|
||||
PoolSize: 50,
|
||||
MaxRetries: maxRetry,
|
||||
})
|
||||
} else {
|
||||
rdb = redis.NewClient(&redis.Options{
|
||||
Addr: redisConf.Address[0],
|
||||
Username: redisConf.Username,
|
||||
Password: redisConf.Password,
|
||||
DB: 0, // use default DB
|
||||
PoolSize: 100, // connection pool size
|
||||
MaxRetries: maxRetry,
|
||||
})
|
||||
}
|
||||
|
||||
var err error
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Second*10)
|
||||
defer cancel()
|
||||
err = rdb.Ping(ctx).Err()
|
||||
if err != nil {
|
||||
errMsg := fmt.Sprintf("address:%s, username:%s, password:%s, clusterMode:%t, enablePipeline:%t", redisConf.Address, redisConf.Username,
|
||||
redisConf.Password, redisConf.ClusterMode, redisConf.EnablePipeline)
|
||||
return nil, errs.WrapMsg(err, errMsg)
|
||||
}
|
||||
redisClient = rdb
|
||||
log.CInfo(ctx, "REDIS connected successfully", "address", redisConf.Address, "username", redisConf.Username, "password", redisConf.Password, "clusterMode", redisConf.ClusterMode, "enablePipeline", redisConf.EnablePipeline)
|
||||
return rdb, err
|
||||
}
|
||||
|
||||
// overrideConfigFromEnv overrides configuration fields with environment variables if present.
|
||||
func overrideConfigFromEnv(redis *config.Redis) {
|
||||
if envAddr := os.Getenv("REDIS_ADDRESS"); envAddr != "" {
|
||||
if envPort := os.Getenv("REDIS_PORT"); envPort != "" {
|
||||
addresses := strings.Split(envAddr, ",")
|
||||
for i, addr := range addresses {
|
||||
addresses[i] = addr + ":" + envPort
|
||||
}
|
||||
redis.Address = addresses
|
||||
} else {
|
||||
redis.Address = strings.Split(envAddr, ",")
|
||||
}
|
||||
}
|
||||
|
||||
if envUser := os.Getenv("REDIS_USERNAME"); envUser != "" {
|
||||
redis.Username = envUser
|
||||
}
|
||||
|
||||
if envPass := os.Getenv("REDIS_PASSWORD"); envPass != "" {
|
||||
redis.Password = envPass
|
||||
}
|
||||
}
|
||||
@ -1 +0,0 @@
|
||||
package unrelation // import "github.com/openimsdk/open-im-server/v3/pkg/common/db/unrelation"
|
||||
@ -1,127 +0,0 @@
|
||||
// Copyright © 2023 OpenIM. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package unrelation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/openimsdk/open-im-server/v3/pkg/common/config"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
"github.com/openimsdk/tools/log"
|
||||
"github.com/openimsdk/tools/mw/specialerror"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
const (
|
||||
maxRetry = 10 // number of retries
|
||||
mongoConnTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
type Mongo struct {
|
||||
db *mongo.Client
|
||||
mongoConf *config.Mongo
|
||||
}
|
||||
|
||||
// NewMongoDB Initialize MongoDB connection.
|
||||
func NewMongoDB(ctx context.Context, mongoConf *config.Mongo) (*Mongo, error) {
|
||||
specialerror.AddReplace(mongo.ErrNoDocuments, errs.ErrRecordNotFound)
|
||||
uri := buildMongoURI(mongoConf)
|
||||
|
||||
var mongoClient *mongo.Client
|
||||
var err error
|
||||
|
||||
// Retry connecting to MongoDB
|
||||
for i := 0; i <= maxRetry; i++ {
|
||||
ctx, cancel := context.WithTimeout(ctx, mongoConnTimeout)
|
||||
defer cancel()
|
||||
mongoClient, err = mongo.Connect(ctx, options.Client().ApplyURI(uri))
|
||||
if err == nil {
|
||||
if err = mongoClient.Ping(ctx, nil); err != nil {
|
||||
return nil, errs.WrapMsg(err, uri)
|
||||
}
|
||||
log.CInfo(ctx, "MONGODB connected successfully", "uri", uri)
|
||||
return &Mongo{db: mongoClient, mongoConf: mongoConf}, nil
|
||||
}
|
||||
if shouldRetry(err) {
|
||||
time.Sleep(time.Second) // exponential backoff could be implemented here
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil, errs.WrapMsg(err, uri)
|
||||
}
|
||||
|
||||
func buildMongoURI(mongoConf *config.Mongo) string {
|
||||
uri := os.Getenv("MONGO_URI")
|
||||
if uri != "" {
|
||||
return uri
|
||||
}
|
||||
|
||||
if mongoConf.Uri != "" {
|
||||
return mongoConf.Uri
|
||||
}
|
||||
|
||||
username := os.Getenv("MONGO_OPENIM_USERNAME")
|
||||
password := os.Getenv("MONGO_OPENIM_PASSWORD")
|
||||
address := os.Getenv("MONGO_ADDRESS")
|
||||
port := os.Getenv("MONGO_PORT")
|
||||
database := os.Getenv("MONGO_DATABASE")
|
||||
maxPoolSize := os.Getenv("MONGO_MAX_POOL_SIZE")
|
||||
|
||||
if username == "" {
|
||||
username = mongoConf.Username
|
||||
}
|
||||
if password == "" {
|
||||
password = mongoConf.Password
|
||||
}
|
||||
if address == "" {
|
||||
address = strings.Join(mongoConf.Address, ",")
|
||||
} else if port != "" {
|
||||
address = fmt.Sprintf("%s:%s", address, port)
|
||||
}
|
||||
if database == "" {
|
||||
database = mongoConf.Database
|
||||
}
|
||||
if maxPoolSize == "" {
|
||||
maxPoolSize = fmt.Sprint(mongoConf.MaxPoolSize)
|
||||
}
|
||||
|
||||
if username != "" && password != "" {
|
||||
|
||||
return fmt.Sprintf("mongodb://%s:%s@%s/%s?maxPoolSize=%s", username, password, address, database, maxPoolSize)
|
||||
}
|
||||
return fmt.Sprintf("mongodb://%s/%s?maxPoolSize=%s", address, database, maxPoolSize)
|
||||
}
|
||||
|
||||
func shouldRetry(err error) bool {
|
||||
if cmdErr, ok := err.(mongo.CommandError); ok {
|
||||
return cmdErr.Code != 13 && cmdErr.Code != 18
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// GetClient returns the MongoDB client.
|
||||
func (m *Mongo) GetClient() *mongo.Client {
|
||||
return m.db
|
||||
}
|
||||
|
||||
// GetDatabase returns the specific database from MongoDB.
|
||||
func (m *Mongo) GetDatabase(database string) *mongo.Database {
|
||||
return m.db.Database(database)
|
||||
}
|
||||
Loading…
Reference in new issue