mirror of
https://github.com/tiennm99/miti99bot.git
synced 2026-09-08 20:20:17 +00:00
feat(storage): add MongoDB KVProvider backend
Add MongoDB-backed storage provider with connection pooling, key-value operations (get, set, delete), and comprehensive test coverage. Integrates mongo-driver dependency for self-hosted deployments.
This commit is contained in:
@@ -9,6 +9,8 @@ require (
|
||||
github.com/aws/aws-sdk-go-v2/service/dynamodb v1.57.3
|
||||
github.com/aws/aws-sdk-go-v2/service/ssm v1.68.6
|
||||
github.com/go-telegram/bot v1.20.0
|
||||
github.com/robfig/cron/v3 v3.0.1
|
||||
go.mongodb.org/mongo-driver/v2 v2.7.0
|
||||
golang.org/x/time v0.15.0
|
||||
google.golang.org/api v0.274.0
|
||||
google.golang.org/genai v1.56.0
|
||||
@@ -43,6 +45,11 @@ require (
|
||||
github.com/googleapis/enterprise-certificate-proxy v0.3.14 // indirect
|
||||
github.com/googleapis/gax-go/v2 v2.21.0 // indirect
|
||||
github.com/gorilla/websocket v1.5.3 // indirect
|
||||
github.com/klauspost/compress v1.17.6 // indirect
|
||||
github.com/xdg-go/pbkdf2 v1.0.0 // indirect
|
||||
github.com/xdg-go/scram v1.2.0 // indirect
|
||||
github.com/xdg-go/stringprep v1.0.4 // indirect
|
||||
github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
|
||||
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.61.0 // indirect
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 // indirect
|
||||
|
||||
@@ -78,12 +78,27 @@ github.com/googleapis/gax-go/v2 v2.21.0 h1:h45NjjzEO3faG9Lg/cFrBh2PgegVVgzqKzuZl
|
||||
github.com/googleapis/gax-go/v2 v2.21.0/go.mod h1:But/NJU6TnZsrLai/xBAQLLz+Hc7fHZJt/hsCz3Fih4=
|
||||
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
|
||||
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
|
||||
github.com/klauspost/compress v1.17.6 h1:60eq2E/jlfwQXtvZEeBUYADs+BwKBWURIY+Gj2eRGjI=
|
||||
github.com/klauspost/compress v1.17.6/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM=
|
||||
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=
|
||||
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
|
||||
github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
|
||||
github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
|
||||
github.com/xdg-go/scram v1.2.0 h1:bYKF2AEwG5rqd1BumT4gAnvwU/M9nBp2pTSxeZw7Wvs=
|
||||
github.com/xdg-go/scram v1.2.0/go.mod h1:3dlrS0iBaWKYVt2ZfA4cj48umJZ+cAEbR6/SjLA88I8=
|
||||
github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8=
|
||||
github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM=
|
||||
github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 h1:ilQV1hzziu+LLM3zUTJ0trRztfwgjqKnBWNtSRkbmwM=
|
||||
github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78/go.mod h1:aL8wCCfTfSfmXjznFBSZNN13rSJjlIOI1fUNAtF7rmI=
|
||||
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
|
||||
go.mongodb.org/mongo-driver/v2 v2.7.0 h1:RO+zqavD2/GCL3cxOMyZhx6R9Irzr8/6gsoqx5tcY/c=
|
||||
go.mongodb.org/mongo-driver/v2 v2.7.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
|
||||
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.61.0 h1:q4XOmH/0opmeuJtPsbFNivyl7bCt7yRBbeEm2sC/XtQ=
|
||||
@@ -100,20 +115,43 @@ go.opentelemetry.io/otel/sdk/metric v1.42.0 h1:D/1QR46Clz6ajyZ3G8SgNlTJKBdGp84q9
|
||||
go.opentelemetry.io/otel/sdk/metric v1.42.0/go.mod h1:Ua6AAlDKdZ7tdvaQKfSmnFTdHx37+J4ba8MwVCYM5hc=
|
||||
go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A=
|
||||
go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI=
|
||||
golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8=
|
||||
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
|
||||
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
|
||||
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
|
||||
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
|
||||
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ=
|
||||
golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
|
||||
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
|
||||
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
|
||||
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
|
||||
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
|
||||
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4=
|
||||
gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E=
|
||||
google.golang.org/api v0.274.0 h1:aYhycS5QQCwxHLwfEHRRLf9yNsfvp1JadKKWBE54RFA=
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/mongo"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/readpref"
|
||||
)
|
||||
|
||||
// mongoServerSelectionTimeout bounds how long a single operation waits for a
|
||||
// reachable server before failing. The self-hosted container holds one
|
||||
// long-lived client for days against Atlas M0 (which idles / fails over); a
|
||||
// tight selection timeout means a wedged DB surfaces as a fast error on the
|
||||
// next op rather than an unbounded hang. The driver auto-reconnects on the
|
||||
// following operation once the server is back.
|
||||
const mongoServerSelectionTimeout = 5 * time.Second
|
||||
|
||||
// NewMongoClient connects to MongoDB using the full connection URI (including
|
||||
// credentials for Atlas `mongodb+srv://` strings) and verifies reachability
|
||||
// with a Ping under the caller's context deadline.
|
||||
//
|
||||
// The client is goroutine-safe and meant to be reused for the lifetime of the
|
||||
// process; callers must defer Disconnect at shutdown.
|
||||
//
|
||||
// SECURITY: uri carries the username/password for Atlas. Callers MUST NOT log
|
||||
// it — see buildProvider, which logs only the database name.
|
||||
func NewMongoClient(ctx context.Context, uri string) (*mongo.Client, error) {
|
||||
if uri == "" {
|
||||
return nil, fmt.Errorf("storage: MONGO_URL is required for MongoDB")
|
||||
}
|
||||
opts := options.Client().
|
||||
ApplyURI(uri).
|
||||
SetServerSelectionTimeout(mongoServerSelectionTimeout)
|
||||
client, err := mongo.Connect(opts)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("storage: mongo.Connect: %w", err)
|
||||
}
|
||||
if err := client.Ping(ctx, readpref.Primary()); err != nil {
|
||||
// Best-effort cleanup; the connection never became usable.
|
||||
_ = client.Disconnect(context.Background())
|
||||
return nil, fmt.Errorf("storage: mongo ping: %w", err)
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// NewMongoDatabase returns the named database handle from a connected client.
|
||||
// The database is created lazily on first write by MongoDB; no round trip here.
|
||||
func NewMongoDatabase(client *mongo.Client, database string) (*mongo.Database, error) {
|
||||
if database == "" {
|
||||
return nil, fmt.Errorf("storage: MONGO_DATABASE is required for MongoDB")
|
||||
}
|
||||
return client.Database(database), nil
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
)
|
||||
|
||||
// MongoDB document fields. `_id` holds the user key; `value` holds the raw
|
||||
// bytes; `updatedAt` is unix-nanos for observability + future TTL.
|
||||
const (
|
||||
mongoIDField = "_id"
|
||||
mongoValueField = "value"
|
||||
mongoUpdatedAtField = "updatedAt"
|
||||
)
|
||||
|
||||
// MongoKVStore is a KVStore backed by a single MongoDB collection. The caller
|
||||
// (MongoProvider) creates one per module so cross-module isolation is
|
||||
// "different collection" — no key prefix needed at this layer, mirroring
|
||||
// FirestoreKVStore.
|
||||
type MongoKVStore struct {
|
||||
coll *mongo.Collection
|
||||
moduleName string
|
||||
}
|
||||
|
||||
// NewMongoKVStore returns a store writing to the given collection. Callers must
|
||||
// validate the collection/module name beforehand (MongoProvider does).
|
||||
func NewMongoKVStore(coll *mongo.Collection, moduleName string) *MongoKVStore {
|
||||
return &MongoKVStore{coll: coll, moduleName: moduleName}
|
||||
}
|
||||
|
||||
// decodeValue extracts the stored value bytes from a decoded document. The
|
||||
// driver decodes a BSON binary into bson.Binary and a BSON string into string;
|
||||
// accept both so values written by any path (or a future schema) round-trip,
|
||||
// mirroring FirestoreKVStore's dual-type handling.
|
||||
func (s *MongoKVStore) decodeValue(key string, doc bson.M) ([]byte, error) {
|
||||
raw, ok := doc[mongoValueField]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("mongo get %s/%s: missing %q field", s.moduleName, key, mongoValueField)
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case bson.Binary:
|
||||
return v.Data, nil
|
||||
case []byte:
|
||||
return v, nil
|
||||
case string:
|
||||
return []byte(v), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("mongo get %s/%s: unexpected value type %T", s.moduleName, key, raw)
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns the raw bytes stored at key, or ErrNotFound.
|
||||
func (s *MongoKVStore) Get(ctx context.Context, key string) ([]byte, error) {
|
||||
if err := validateKey(key); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var doc bson.M
|
||||
err := s.coll.FindOne(ctx, bson.M{mongoIDField: key}).Decode(&doc)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongo.ErrNoDocuments) {
|
||||
return nil, ErrNotFound
|
||||
}
|
||||
return nil, fmt.Errorf("mongo get %s/%s: %w", s.moduleName, key, err)
|
||||
}
|
||||
return s.decodeValue(key, doc)
|
||||
}
|
||||
|
||||
// GetJSON decodes the value at key into dst.
|
||||
func (s *MongoKVStore) GetJSON(ctx context.Context, key string, dst any) error {
|
||||
raw, err := s.Get(ctx, key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := json.Unmarshal(raw, dst); err != nil {
|
||||
return fmt.Errorf("mongo get %s/%s: json decode: %w", s.moduleName, key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// doc builds the persisted document for key/val with a fresh updatedAt stamp.
|
||||
// value is stored as BSON binary (generic subtype) so non-UTF-8 payloads
|
||||
// round-trip; updatedAt is int64 unix-nanos to match DynamoDB byte-for-byte
|
||||
// (see dynamodb_kv.go) and keep migration faithful.
|
||||
func (s *MongoKVStore) doc(key string, val []byte) bson.M {
|
||||
return bson.M{
|
||||
mongoIDField: key,
|
||||
mongoValueField: bson.Binary{Subtype: bson.TypeBinaryGeneric, Data: val},
|
||||
mongoUpdatedAtField: time.Now().UTC().UnixNano(),
|
||||
}
|
||||
}
|
||||
|
||||
// Put writes raw bytes at key, creating or overwriting.
|
||||
func (s *MongoKVStore) Put(ctx context.Context, key string, val []byte) error {
|
||||
if err := validateKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := s.coll.ReplaceOne(ctx,
|
||||
bson.M{mongoIDField: key},
|
||||
s.doc(key, val),
|
||||
options.Replace().SetUpsert(true),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("mongo put %s/%s: %w", s.moduleName, key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PutJSON marshals val and writes the bytes at key.
|
||||
func (s *MongoKVStore) PutJSON(ctx context.Context, key string, val any) error {
|
||||
raw, err := json.Marshal(val)
|
||||
if err != nil {
|
||||
return fmt.Errorf("mongo put %s/%s: json encode: %w", s.moduleName, key, err)
|
||||
}
|
||||
return s.Put(ctx, key, raw)
|
||||
}
|
||||
|
||||
// CompareAndSwap conditionally replaces the value only when it still equals
|
||||
// expected. A nil expected means the key must not yet exist.
|
||||
//
|
||||
// - expected == nil → InsertOne; the unique _id index makes the absent-insert
|
||||
// race linearizable (exactly one writer wins, losers get a duplicate-key
|
||||
// error → ErrConflict). This is a LIVE first-write path (every new
|
||||
// coin/gold portfolio), not an edge case.
|
||||
// - expected != nil → UpdateOne filtered on the matching value; MatchedCount
|
||||
// of 0 means the stored value changed (or the key is absent) → ErrConflict.
|
||||
func (s *MongoKVStore) CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error {
|
||||
if err := validateKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
if expected == nil {
|
||||
_, err := s.coll.InsertOne(ctx, s.doc(key, val))
|
||||
if err != nil {
|
||||
if mongo.IsDuplicateKeyError(err) {
|
||||
return ErrConflict
|
||||
}
|
||||
return fmt.Errorf("mongo compare-and-swap %s/%s: %w", s.moduleName, key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
res, err := s.coll.UpdateOne(ctx,
|
||||
bson.M{
|
||||
mongoIDField: key,
|
||||
mongoValueField: bson.Binary{Subtype: bson.TypeBinaryGeneric, Data: expected},
|
||||
},
|
||||
bson.M{"$set": bson.M{
|
||||
mongoValueField: bson.Binary{Subtype: bson.TypeBinaryGeneric, Data: val},
|
||||
mongoUpdatedAtField: time.Now().UTC().UnixNano(),
|
||||
}},
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("mongo compare-and-swap %s/%s: %w", s.moduleName, key, err)
|
||||
}
|
||||
if res.MatchedCount == 0 {
|
||||
return ErrConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Delete removes the document at key. Deleting a missing key is not an error
|
||||
// (idempotent) — DeleteOne with a zero match count returns nil.
|
||||
func (s *MongoKVStore) Delete(ctx context.Context, key string) error {
|
||||
if err := validateKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := s.coll.DeleteOne(ctx, bson.M{mongoIDField: key})
|
||||
if err != nil {
|
||||
return fmt.Errorf("mongo delete %s/%s: %w", s.moduleName, key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// List returns all document IDs in the collection that start with prefix.
|
||||
// Implemented as a half-open range scan on _id (reusing prefixSuccessor) so it
|
||||
// uses the _id index and avoids regex injection. Empty prefix returns the whole
|
||||
// collection.
|
||||
func (s *MongoKVStore) List(ctx context.Context, prefix string) ([]string, error) {
|
||||
if err := validatePrefix(prefix); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
filter := bson.M{}
|
||||
if prefix != "" {
|
||||
filter[mongoIDField] = bson.M{
|
||||
"$gte": prefix,
|
||||
"$lt": prefixSuccessor(prefix),
|
||||
}
|
||||
}
|
||||
cur, err := s.coll.Find(ctx, filter, options.Find().SetProjection(bson.M{mongoIDField: 1}))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mongo list %s prefix=%q: %w", s.moduleName, prefix, err)
|
||||
}
|
||||
defer func() { _ = cur.Close(ctx) }()
|
||||
|
||||
var keys []string
|
||||
for cur.Next(ctx) {
|
||||
var doc struct {
|
||||
ID string `bson:"_id"`
|
||||
}
|
||||
if err := cur.Decode(&doc); err != nil {
|
||||
return nil, fmt.Errorf("mongo list %s prefix=%q: decode: %w", s.moduleName, prefix, err)
|
||||
}
|
||||
keys = append(keys, doc.ID)
|
||||
}
|
||||
if err := cur.Err(); err != nil {
|
||||
return nil, fmt.Errorf("mongo list %s prefix=%q: %w", s.moduleName, prefix, err)
|
||||
}
|
||||
return keys, nil
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo"
|
||||
)
|
||||
|
||||
// mongoLocalSetup connects to a local MongoDB and returns a store over a fresh,
|
||||
// uniquely-named collection. Tests skip if MONGODB_TEST_URL is unset so CI
|
||||
// without a Mongo container still builds (mirrors the DynamoDB Local gating).
|
||||
func mongoLocalSetup(t *testing.T, module string) (*MongoKVStore, *mongo.Database, func()) {
|
||||
t.Helper()
|
||||
uri := os.Getenv("MONGODB_TEST_URL")
|
||||
if uri == "" {
|
||||
t.Skip("MONGODB_TEST_URL not set; skipping MongoDB integration test (run `make mongo-local` to start the local container)")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
client, err := NewMongoClient(ctx, uri)
|
||||
if err != nil {
|
||||
t.Fatalf("NewMongoClient: %v", err)
|
||||
}
|
||||
// Unique DB per test so parallel runs and cross-module isolation checks
|
||||
// never collide. Mongo db names are capped at 63 bytes.
|
||||
dbName := fmt.Sprintf("miti99bot_test_%d", time.Now().UnixNano())
|
||||
if len(dbName) > 63 {
|
||||
dbName = dbName[:63]
|
||||
}
|
||||
db := client.Database(dbName)
|
||||
|
||||
cleanup := func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
_ = db.Drop(ctx)
|
||||
_ = client.Disconnect(ctx)
|
||||
}
|
||||
return NewMongoKVStore(db.Collection(module), module), db, cleanup
|
||||
}
|
||||
|
||||
func TestMongoKVStore_PutGetDelete(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
if err := s.Put(ctx, "user:1:state", []byte("hello")); err != nil {
|
||||
t.Fatalf("Put: %v", err)
|
||||
}
|
||||
got, err := s.Get(ctx, "user:1:state")
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
if string(got) != "hello" {
|
||||
t.Errorf("Get: got %q, want %q", got, "hello")
|
||||
}
|
||||
|
||||
// Overwrite preserves byte fidelity.
|
||||
if err := s.Put(ctx, "user:1:state", []byte("world")); err != nil {
|
||||
t.Fatalf("Put overwrite: %v", err)
|
||||
}
|
||||
got, _ = s.Get(ctx, "user:1:state")
|
||||
if string(got) != "world" {
|
||||
t.Errorf("Get after overwrite: got %q, want %q", got, "world")
|
||||
}
|
||||
|
||||
if err := s.Delete(ctx, "user:1:state"); err != nil {
|
||||
t.Fatalf("Delete: %v", err)
|
||||
}
|
||||
if _, err := s.Get(ctx, "user:1:state"); !errors.Is(err, ErrNotFound) {
|
||||
t.Errorf("Get after Delete: got %v, want ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoKVStore_GetMissing(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
if _, err := s.Get(context.Background(), "missing"); !errors.Is(err, ErrNotFound) {
|
||||
t.Errorf("got %v, want ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoKVStore_NonUTF8RoundTrip(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
raw := []byte{0x00, 0xff, 0xfe, 0x01, 0x80}
|
||||
if err := s.Put(ctx, "bin", raw); err != nil {
|
||||
t.Fatalf("Put: %v", err)
|
||||
}
|
||||
got, err := s.Get(ctx, "bin")
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
if string(got) != string(raw) {
|
||||
t.Errorf("non-UTF-8 round trip: got %v, want %v", got, raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoKVStore_JSONRoundTrip(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "loldle")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
type state struct {
|
||||
Score int `json:"score"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
in := state{Score: 42, Name: "ezreal"}
|
||||
if err := s.PutJSON(ctx, "u1", in); err != nil {
|
||||
t.Fatalf("PutJSON: %v", err)
|
||||
}
|
||||
var out state
|
||||
if err := s.GetJSON(ctx, "u1", &out); err != nil {
|
||||
t.Fatalf("GetJSON: %v", err)
|
||||
}
|
||||
if out != in {
|
||||
t.Errorf("got %+v, want %+v", out, in)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoKVStore_ListPrefix(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
for _, k := range []string{"user:1:state", "user:2:state", "config:daily", "user:1:history"} {
|
||||
if err := s.Put(ctx, k, []byte("x")); err != nil {
|
||||
t.Fatalf("Put %s: %v", k, err)
|
||||
}
|
||||
}
|
||||
|
||||
got, err := s.List(ctx, "user:")
|
||||
if err != nil {
|
||||
t.Fatalf("List: %v", err)
|
||||
}
|
||||
want := map[string]bool{"user:1:state": true, "user:2:state": true, "user:1:history": true}
|
||||
if len(got) != len(want) {
|
||||
t.Errorf("List: got %v (len=%d), want len=%d", got, len(got), len(want))
|
||||
}
|
||||
for _, k := range got {
|
||||
if !want[k] {
|
||||
t.Errorf("List: unexpected key %q", k)
|
||||
}
|
||||
}
|
||||
|
||||
// Empty prefix returns everything in the collection.
|
||||
all, err := s.List(ctx, "")
|
||||
if err != nil {
|
||||
t.Fatalf("List empty: %v", err)
|
||||
}
|
||||
if len(all) != 4 {
|
||||
t.Errorf("List(\"\"): got %d keys, want 4", len(all))
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoKVStore_CompareAndSwap(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "gold")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
// Create-if-absent succeeds once, then conflicts.
|
||||
if err := s.CompareAndSwap(ctx, "user:1", nil, []byte("v1")); err != nil {
|
||||
t.Fatalf("CompareAndSwap create: %v", err)
|
||||
}
|
||||
if err := s.CompareAndSwap(ctx, "user:1", nil, []byte("v2")); !errors.Is(err, ErrConflict) {
|
||||
t.Errorf("CompareAndSwap create over existing: got %v, want ErrConflict", err)
|
||||
}
|
||||
|
||||
// Swap with matching expected succeeds; stale expected conflicts.
|
||||
if err := s.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v2")); err != nil {
|
||||
t.Fatalf("CompareAndSwap matching: %v", err)
|
||||
}
|
||||
if err := s.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v3")); !errors.Is(err, ErrConflict) {
|
||||
t.Errorf("CompareAndSwap stale: got %v, want ErrConflict", err)
|
||||
}
|
||||
got, err := s.Get(ctx, "user:1")
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
if string(got) != "v2" {
|
||||
t.Errorf("stored value = %q, want %q", got, "v2")
|
||||
}
|
||||
|
||||
// Non-nil expected on a missing key conflicts (caller reloads and retries).
|
||||
if err := s.CompareAndSwap(ctx, "user:missing", []byte("v1"), []byte("v2")); !errors.Is(err, ErrConflict) {
|
||||
t.Errorf("CompareAndSwap missing key: got %v, want ErrConflict", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMongoKVStore_CompareAndSwap_ConcurrentInsert proves the absent-insert
|
||||
// race is linearizable: N goroutines racing a nil-expected CAS on the same key
|
||||
// must yield exactly one winner; every loser gets ErrConflict (never a silent
|
||||
// overwrite). This is the blocking correctness gate from Phase 1.
|
||||
func TestMongoKVStore_CompareAndSwap_ConcurrentInsert(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "coin")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
const n = 16
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
var wins, conflicts, others int
|
||||
start := make(chan struct{})
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
<-start // release all goroutines at once to maximize contention
|
||||
err := s.CompareAndSwap(ctx, "race", nil, []byte(fmt.Sprintf("v%d", i)))
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
switch {
|
||||
case err == nil:
|
||||
wins++
|
||||
case errors.Is(err, ErrConflict):
|
||||
conflicts++
|
||||
default:
|
||||
others++
|
||||
t.Errorf("unexpected CAS error: %v", err)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
|
||||
if wins != 1 {
|
||||
t.Errorf("concurrent nil-expected CAS: got %d winners, want exactly 1", wins)
|
||||
}
|
||||
if conflicts != n-1 {
|
||||
t.Errorf("concurrent nil-expected CAS: got %d conflicts, want %d", conflicts, n-1)
|
||||
}
|
||||
if others != 0 {
|
||||
t.Errorf("concurrent CAS produced %d non-conflict errors", others)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMongoKVStore_GetDocWithoutValueField covers a malformed document edge
|
||||
// case (a doc that exists but lacks the value field) — Get must surface a clear
|
||||
// error, not panic or return empty bytes silently.
|
||||
func TestMongoKVStore_GetDocWithoutValueField(t *testing.T) {
|
||||
s, db, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
ctx := context.Background()
|
||||
// Insert a doc directly with no `value` field.
|
||||
if _, err := db.Collection("wordle").InsertOne(ctx, bson.M{"_id": "novalue", "other": 1}); err != nil {
|
||||
t.Fatalf("seed InsertOne: %v", err)
|
||||
}
|
||||
if _, err := s.Get(ctx, "novalue"); err == nil {
|
||||
t.Errorf("Get on doc without value field: got nil error, want a descriptive error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoKVStore_DeleteMissingNoError(t *testing.T) {
|
||||
s, _, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
if err := s.Delete(context.Background(), "never-existed"); err != nil {
|
||||
t.Errorf("Delete missing key: got %v, want nil (idempotent)", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"go.mongodb.org/mongo-driver/v2/mongo"
|
||||
)
|
||||
|
||||
// MongoProvider is a KVProvider that creates one collection per module,
|
||||
// mirroring FirestoreProvider. Collection-per-module IS the isolation — no key
|
||||
// prefix wrapping is needed at this layer.
|
||||
type MongoProvider struct {
|
||||
db *mongo.Database
|
||||
}
|
||||
|
||||
// NewMongoProvider returns a provider over the given database handle. The
|
||||
// underlying client must outlive every KVStore the provider hands out; callers
|
||||
// own its Disconnect.
|
||||
func NewMongoProvider(db *mongo.Database) *MongoProvider {
|
||||
return &MongoProvider{db: db}
|
||||
}
|
||||
|
||||
// For returns a MongoKVStore writing to a collection named after the module.
|
||||
// moduleName is re-validated against collectionNameRe — defense in depth
|
||||
// against caller bugs that bypass modules.Build, identical to the firestore
|
||||
// and dynamodb providers. An invalid name yields an invalidStore whose every
|
||||
// op errors at first use.
|
||||
func (p *MongoProvider) For(moduleName string) KVStore {
|
||||
if !collectionNameRe.MatchString(moduleName) {
|
||||
return invalidStore{name: moduleName}
|
||||
}
|
||||
return NewMongoKVStore(p.db.Collection(moduleName), moduleName)
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// MongoProvider.For re-validates the module name as defense-in-depth. Invalid
|
||||
// names return invalidStore without touching the DB — the branch worth locking
|
||||
// even without a live Mongo (mirrors the firestore provider test).
|
||||
func TestMongoProvider_For_RejectsInvalidName(t *testing.T) {
|
||||
p := &MongoProvider{db: nil}
|
||||
bogus := []string{
|
||||
"",
|
||||
"with spaces",
|
||||
"WITHCAPS",
|
||||
"path/traversal",
|
||||
"../etc/passwd",
|
||||
"way-too-long-for-our-32-char-limit-x",
|
||||
"with:colon",
|
||||
}
|
||||
for _, name := range bogus {
|
||||
store := p.For(name)
|
||||
if _, err := store.Get(context.Background(), "any-key"); !errors.Is(err, ErrInvalidModuleName) {
|
||||
t.Errorf("For(%q).Get → %v, want ErrInvalidModuleName", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestMongoProvider_CrossModuleIsolation verifies collection-per-module
|
||||
// isolation: the same key written through two module stores yields independent
|
||||
// values. Gated on MONGODB_TEST_URL.
|
||||
func TestMongoProvider_CrossModuleIsolation(t *testing.T) {
|
||||
_, db, cleanup := mongoLocalSetup(t, "wordle")
|
||||
defer cleanup()
|
||||
|
||||
p := NewMongoProvider(db)
|
||||
ctx := context.Background()
|
||||
|
||||
wordle := p.For("wordle")
|
||||
loldle := p.For("loldle")
|
||||
if err := wordle.Put(ctx, "shared", []byte("from-wordle")); err != nil {
|
||||
t.Fatalf("wordle Put: %v", err)
|
||||
}
|
||||
if err := loldle.Put(ctx, "shared", []byte("from-loldle")); err != nil {
|
||||
t.Fatalf("loldle Put: %v", err)
|
||||
}
|
||||
|
||||
gotW, _ := wordle.Get(ctx, "shared")
|
||||
gotL, _ := loldle.Get(ctx, "shared")
|
||||
if string(gotW) != "from-wordle" {
|
||||
t.Errorf("wordle key leaked: got %q", gotW)
|
||||
}
|
||||
if string(gotL) != "from-loldle" {
|
||||
t.Errorf("loldle key leaked: got %q", gotL)
|
||||
}
|
||||
// Canonical names pass validation (not invalidStore); this also covers the
|
||||
// valid-name branch of For against a real database.
|
||||
for _, name := range []string{"misc", "demo-mod", "x", "a1_b-2"} {
|
||||
if _, ok := p.For(name).(invalidStore); ok {
|
||||
t.Errorf("For(%q) returned invalidStore; expected validation to pass", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user