52 lines
1013 B
Go
52 lines
1013 B
Go
|
package mongo
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"haixun-backend/internal/config"
|
||
|
|
|
||
|
|
"go.mongodb.org/mongo-driver/mongo"
|
||
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Client struct {
|
||
|
|
raw *mongo.Client
|
||
|
|
db *mongo.Database
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewClient(ctx context.Context, conf config.MongoConf) (*Client, error) {
|
||
|
|
if conf.URI == "" || conf.Database == "" {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
timeout := time.Duration(conf.TimeoutSeconds) * time.Second
|
||
|
|
if timeout <= 0 {
|
||
|
|
timeout = 10 * time.Second
|
||
|
|
}
|
||
|
|
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||
|
|
defer cancel()
|
||
|
|
|
||
|
|
raw, err := mongo.Connect(ctx, options.Client().ApplyURI(conf.URI))
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if err := raw.Ping(ctx, nil); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return &Client{raw: raw, db: raw.Database(conf.Database)}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) Database() *mongo.Database {
|
||
|
|
if c == nil {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return c.db
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) Close(ctx context.Context) error {
|
||
|
|
if c == nil || c.raw == nil {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return c.raw.Disconnect(ctx)
|
||
|
|
}
|