From bfaecef86c904ecdc01d7a5647870b7264a32255 Mon Sep 17 00:00:00 2001 From: tiennm99 Date: Sun, 28 Jun 2026 09:58:07 +0700 Subject: [PATCH] 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. --- go.mod | 7 + go.sum | 38 +++ internal/storage/mongodb_client.go | 56 +++++ internal/storage/mongodb_kv.go | 214 +++++++++++++++++ internal/storage/mongodb_kv_test.go | 272 ++++++++++++++++++++++ internal/storage/mongodb_provider.go | 31 +++ internal/storage/mongodb_provider_test.go | 65 ++++++ 7 files changed, 683 insertions(+) create mode 100644 internal/storage/mongodb_client.go create mode 100644 internal/storage/mongodb_kv.go create mode 100644 internal/storage/mongodb_kv_test.go create mode 100644 internal/storage/mongodb_provider.go create mode 100644 internal/storage/mongodb_provider_test.go diff --git a/go.mod b/go.mod index a694a14..b374456 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 232fbad..71eac91 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/internal/storage/mongodb_client.go b/internal/storage/mongodb_client.go new file mode 100644 index 0000000..d76b0c1 --- /dev/null +++ b/internal/storage/mongodb_client.go @@ -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 +} diff --git a/internal/storage/mongodb_kv.go b/internal/storage/mongodb_kv.go new file mode 100644 index 0000000..653495d --- /dev/null +++ b/internal/storage/mongodb_kv.go @@ -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 +} diff --git a/internal/storage/mongodb_kv_test.go b/internal/storage/mongodb_kv_test.go new file mode 100644 index 0000000..ab7e177 --- /dev/null +++ b/internal/storage/mongodb_kv_test.go @@ -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) + } +} diff --git a/internal/storage/mongodb_provider.go b/internal/storage/mongodb_provider.go new file mode 100644 index 0000000..785c708 --- /dev/null +++ b/internal/storage/mongodb_provider.go @@ -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) +} diff --git a/internal/storage/mongodb_provider_test.go b/internal/storage/mongodb_provider_test.go new file mode 100644 index 0000000..1f6dcbe --- /dev/null +++ b/internal/storage/mongodb_provider_test.go @@ -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) + } + } +}