mirror of
https://github.com/tiennm99/keepalive.git
synced 2026-10-11 03:13:31 +00:00
tiennm99
47afb85093
- Generated service names come only from a URL host, a key=value DSN's host=, or a MySQL tcp() address, so a password in a key=value DSN can no longer end up in the name printed on every log line. - connect_timeout is added only to postgres:// and postgresql:// URLs (parsed, so it never lands after a fragment) or as a key=value token, never glued onto a key=value value containing "://". - Driver parse errors that quote password fragments (mongo escape errors, lib/pq's missing "=" error) are replaced with generic hints. - MongoDB keeps its client only after a successful connect, so a failed connect is not disconnected twice. - Couchbase gets ready_timeout plus 1 minute to connect, so raising ready_timeout takes effect. - Shutdown waits at most 7 seconds, inside Docker's 10-second grace period. - Each adapter declares its config keys; unknown keys anywhere in the file, including under config, log a warning without blocking start.
90 lines
2.2 KiB
Go
90 lines
2.2 KiB
Go
package adapter
|
|
|
|
import (
|
|
"context"
|
|
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"go.mongodb.org/mongo-driver/v2/mongo"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
)
|
|
|
|
func init() {
|
|
ConfigKeys["mongodb"] = []string{"uri", "database", "collection"}
|
|
ConfigKeys["mongo"] = ConfigKeys["mongodb"]
|
|
Registry["mongodb"] = func(cfg Config) (Adapter, error) {
|
|
uri, err := cfg.Required("uri")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
dbName, err := cfg.Required("database")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
collName, err := cfg.Required("collection")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &mongoAdapter{
|
|
uri: uri,
|
|
dbName: dbName,
|
|
collName: collName,
|
|
docID: cfg.Optional("counter_key", "counter"),
|
|
}, nil
|
|
}
|
|
Registry["mongo"] = Registry["mongodb"]
|
|
}
|
|
|
|
type mongoAdapter struct {
|
|
client *mongo.Client
|
|
coll *mongo.Collection
|
|
uri string
|
|
dbName string
|
|
collName string
|
|
docID string
|
|
}
|
|
|
|
func (a *mongoAdapter) Connect(ctx context.Context) error {
|
|
client, err := mongo.Connect(options.Client().ApplyURI(a.uri))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
a.coll = client.Database(a.dbName).Collection(a.collName)
|
|
if err := client.Ping(ctx, nil); err != nil {
|
|
client.Disconnect(ctx)
|
|
return err
|
|
}
|
|
if err := a.ensureInitialized(ctx); err != nil {
|
|
client.Disconnect(ctx)
|
|
return err
|
|
}
|
|
// Set only on success so Close does not disconnect a failed client again.
|
|
a.client = client
|
|
return nil
|
|
}
|
|
|
|
func (a *mongoAdapter) ensureInitialized(ctx context.Context) error {
|
|
update := bson.M{"$setOnInsert": bson.M{"count": int64(0)}}
|
|
_, err := a.coll.UpdateOne(ctx, bson.M{"_id": a.docID}, update, options.UpdateOne().SetUpsert(true))
|
|
return err
|
|
}
|
|
|
|
func (a *mongoAdapter) Increment(ctx context.Context) (int64, error) {
|
|
update := bson.M{"$inc": bson.M{"count": 1}}
|
|
opts := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After)
|
|
var out struct {
|
|
Count int64 `bson:"count"`
|
|
}
|
|
err := a.coll.FindOneAndUpdate(ctx, bson.M{"_id": a.docID}, update, opts).Decode(&out)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return out.Count, nil
|
|
}
|
|
|
|
func (a *mongoAdapter) Close(ctx context.Context) error {
|
|
if a.client == nil {
|
|
return nil
|
|
}
|
|
return a.client.Disconnect(ctx)
|
|
}
|