Skip to content

Commit 9938f7d

Browse files
h3n4lclaude
andauthored
feat: support db.coll.count() and cursor count/itcount/size (#28)
* feat: support db.coll.count() and cursor count/itcount/size Adds the deprecated-but-heavily-used count() forms by routing through the modern driver methods: - db.coll.count() (zero args) -> EstimatedDocumentCount, preserving the fast metadata path mongosh users expect on large collections. - db.coll.count(filter, opts?) -> CountDocuments, the only correct option once a filter is involved (EstimatedDocumentCount cannot take a filter). - cursor.count() / itcount() / size() on a find cursor -> CountDocuments with the accumulated filter/skip/limit/hint. Modern mongosh always honors skip/limit on cursor.count, so we do too. Aggregate-cursor itcount() requires pipeline rewriting ($count stage) and is left unsupported until telemetry shows demand. Motivated by gomongo-fallback telemetry from production: count() forms account for the largest single bucket of mongosh fallbacks. Removing them is a prerequisite for retiring the mongosh fallback path. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * review: reject cursor count/itcount/size args; tighten unsupported assert Address Copilot review feedback on PR #28. - translate.go: cursor count/itcount/size now reject any positional arguments rather than silently dropping them. mongosh's legacy cursor.count(applySkipLimit) boolean is a no-op in modern drivers (skip/limit always apply), so accepting it would make incorrect caller assumptions hard to debug. Updated comment to spell out *why* the OpType is retargeted instead of running find then count on the cursor (mongosh issues a separate count command, never iterates). - collection_test.go: TestAggregateItcountUnsupported now asserts the Operation field on the error so a different unsupported method can't slip past as a false positive. - Adds TestCursorCountRejectsArgs covering count(true)/count(false)/ itcount(true)/size(1). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent f6c50d4 commit 9938f7d

2 files changed

Lines changed: 211 additions & 4 deletions

File tree

collection_test.go

Lines changed: 180 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1943,6 +1943,111 @@ func TestEstimatedDocumentCountMaxTimeMS(t *testing.T) {
19431943
})
19441944
}
19451945

1946+
// TestCollectionCount covers db.coll.count() — the deprecated collection method
1947+
// still heavily used in mongosh scripts. Zero-arg routes to estimatedDocumentCount
1948+
// (preserving the fast metadata path); any-arg routes to countDocuments.
1949+
func TestCollectionCount(t *testing.T) {
1950+
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
1951+
dbName := fmt.Sprintf("testdb_coll_count_%s", db.Name)
1952+
defer testutil.CleanupDatabase(t, db.Client, dbName)
1953+
1954+
ctx := context.Background()
1955+
1956+
coll := db.Client.Database(dbName).Collection("users")
1957+
_, err := coll.InsertMany(ctx, []any{
1958+
bson.M{"name": "alice", "age": 30, "status": "active"},
1959+
bson.M{"name": "bob", "age": 25, "status": "inactive"},
1960+
bson.M{"name": "charlie", "age": 35, "status": "active"},
1961+
bson.M{"name": "diana", "age": 28, "status": "active"},
1962+
})
1963+
require.NoError(t, err)
1964+
1965+
gc := gomongo.NewClient(db.Client)
1966+
1967+
// count() — zero args, fast path
1968+
result, err := gc.Execute(ctx, dbName, "db.users.count()")
1969+
require.NoError(t, err)
1970+
require.Equal(t, int64(4), result.Value[0].(int64))
1971+
1972+
// count({}) — explicit empty filter, accurate path
1973+
result, err = gc.Execute(ctx, dbName, "db.users.count({})")
1974+
require.NoError(t, err)
1975+
require.Equal(t, int64(4), result.Value[0].(int64))
1976+
1977+
// count(filter)
1978+
result, err = gc.Execute(ctx, dbName, `db.users.count({ status: "active" })`)
1979+
require.NoError(t, err)
1980+
require.Equal(t, int64(3), result.Value[0].(int64))
1981+
1982+
// count(filter) with comparison operator
1983+
result, err = gc.Execute(ctx, dbName, `db.users.count({ age: { $gte: 30 } })`)
1984+
require.NoError(t, err)
1985+
require.Equal(t, int64(2), result.Value[0].(int64))
1986+
})
1987+
}
1988+
1989+
func TestCollectionCountWithOptions(t *testing.T) {
1990+
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
1991+
dbName := fmt.Sprintf("testdb_coll_count_opts_%s", db.Name)
1992+
defer testutil.CleanupDatabase(t, db.Client, dbName)
1993+
1994+
ctx := context.Background()
1995+
1996+
coll := db.Client.Database(dbName).Collection("users")
1997+
_, err := coll.InsertMany(ctx, []any{
1998+
bson.M{"name": "alice", "age": 30},
1999+
bson.M{"name": "bob", "age": 25},
2000+
bson.M{"name": "charlie", "age": 35},
2001+
bson.M{"name": "diana", "age": 28},
2002+
bson.M{"name": "eve", "age": 40},
2003+
})
2004+
require.NoError(t, err)
2005+
2006+
gc := gomongo.NewClient(db.Client)
2007+
2008+
// limit option
2009+
result, err := gc.Execute(ctx, dbName, `db.users.count({}, { limit: 3 })`)
2010+
require.NoError(t, err)
2011+
require.Equal(t, int64(3), result.Value[0].(int64))
2012+
2013+
// skip option
2014+
result, err = gc.Execute(ctx, dbName, `db.users.count({}, { skip: 2 })`)
2015+
require.NoError(t, err)
2016+
require.Equal(t, int64(3), result.Value[0].(int64))
2017+
2018+
// skip + limit
2019+
result, err = gc.Execute(ctx, dbName, `db.users.count({}, { skip: 1, limit: 2 })`)
2020+
require.NoError(t, err)
2021+
require.Equal(t, int64(2), result.Value[0].(int64))
2022+
2023+
// maxTimeMS option
2024+
result, err = gc.Execute(ctx, dbName, `db.users.count({}, { maxTimeMS: 5000 })`)
2025+
require.NoError(t, err)
2026+
require.Equal(t, int64(5), result.Value[0].(int64))
2027+
})
2028+
}
2029+
2030+
func TestCollectionCountEmptyCollection(t *testing.T) {
2031+
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
2032+
dbName := fmt.Sprintf("testdb_coll_count_empty_%s", db.Name)
2033+
defer testutil.CleanupDatabase(t, db.Client, dbName)
2034+
2035+
ctx := context.Background()
2036+
2037+
gc := gomongo.NewClient(db.Client)
2038+
2039+
// Zero-arg count on missing collection
2040+
result, err := gc.Execute(ctx, dbName, "db.users.count()")
2041+
require.NoError(t, err)
2042+
require.Equal(t, int64(0), result.Value[0].(int64))
2043+
2044+
// count({}) on missing collection
2045+
result, err = gc.Execute(ctx, dbName, "db.users.count({})")
2046+
require.NoError(t, err)
2047+
require.Equal(t, int64(0), result.Value[0].(int64))
2048+
})
2049+
}
2050+
19462051
func TestDistinct(t *testing.T) {
19472052
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
19482053
dbName := fmt.Sprintf("testdb_distinct_%s", db.Name)
@@ -2116,22 +2221,93 @@ func TestDistinctMaxTimeMS(t *testing.T) {
21162221
})
21172222
}
21182223

2119-
func TestCursorCountUnsupported(t *testing.T) {
2224+
func TestCursorCount(t *testing.T) {
21202225
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
21212226
dbName := fmt.Sprintf("testdb_cursor_count_%s", db.Name)
21222227
defer testutil.CleanupDatabase(t, db.Client, dbName)
21232228

21242229
ctx := context.Background()
21252230

2231+
coll := db.Client.Database(dbName).Collection("users")
2232+
_, err := coll.InsertMany(ctx, []any{
2233+
bson.M{"name": "Alice", "status": "active"},
2234+
bson.M{"name": "Bob", "status": "inactive"},
2235+
bson.M{"name": "Charlie", "status": "active"},
2236+
bson.M{"name": "Diana", "status": "active"},
2237+
})
2238+
require.NoError(t, err)
2239+
2240+
gc := gomongo.NewClient(db.Client)
2241+
2242+
// find().count() — empty filter
2243+
result, err := gc.Execute(ctx, dbName, "db.users.find().count()")
2244+
require.NoError(t, err)
2245+
require.Equal(t, int64(4), result.Value[0].(int64))
2246+
2247+
// find(filter).count() — accumulates filter from find()
2248+
result, err = gc.Execute(ctx, dbName, `db.users.find({ status: "active" }).count()`)
2249+
require.NoError(t, err)
2250+
require.Equal(t, int64(3), result.Value[0].(int64))
2251+
2252+
// find().skip().limit().count() — modern mongosh always honors skip/limit on cursor.count
2253+
result, err = gc.Execute(ctx, dbName, "db.users.find().skip(1).limit(2).count()")
2254+
require.NoError(t, err)
2255+
require.Equal(t, int64(2), result.Value[0].(int64))
2256+
2257+
// itcount and size are aliases for count on a find cursor
2258+
result, err = gc.Execute(ctx, dbName, `db.users.find({ status: "active" }).itcount()`)
2259+
require.NoError(t, err)
2260+
require.Equal(t, int64(3), result.Value[0].(int64))
2261+
2262+
result, err = gc.Execute(ctx, dbName, `db.users.find({ status: "active" }).size()`)
2263+
require.NoError(t, err)
2264+
require.Equal(t, int64(3), result.Value[0].(int64))
2265+
})
2266+
}
2267+
2268+
func TestAggregateItcountUnsupported(t *testing.T) {
2269+
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
2270+
dbName := fmt.Sprintf("testdb_agg_itcount_%s", db.Name)
2271+
defer testutil.CleanupDatabase(t, db.Client, dbName)
2272+
2273+
ctx := context.Background()
2274+
21262275
gc := gomongo.NewClient(db.Client)
21272276

2128-
// cursor.count() is not in the planned registry, should return UnsupportedOperationError
2129-
_, err := gc.Execute(ctx, dbName, "db.users.find().count()")
2277+
// aggregate cursors expose itcount() in mongosh but require pipeline rewriting
2278+
// ($count stage). Out of scope until telemetry shows demand.
2279+
_, err := gc.Execute(ctx, dbName, "db.users.aggregate([{ $match: {} }]).itcount()")
21302280
require.Error(t, err)
21312281

21322282
var unsupportedErr *gomongo.UnsupportedOperationError
21332283
require.ErrorAs(t, err, &unsupportedErr)
2134-
require.Equal(t, "count()", unsupportedErr.Operation)
2284+
require.Equal(t, "itcount()", unsupportedErr.Operation)
2285+
})
2286+
}
2287+
2288+
// TestCursorCountRejectsArgs guards against mongosh's legacy
2289+
// cursor.count(applySkipLimit) form silently dropping the boolean: the arg is
2290+
// a no-op in modern drivers, and silently dropping it would make incorrect
2291+
// caller assumptions hard to debug.
2292+
func TestCursorCountRejectsArgs(t *testing.T) {
2293+
testutil.RunOnAllDBs(t, func(t *testing.T, db testutil.TestDB) {
2294+
dbName := fmt.Sprintf("testdb_cursor_count_args_%s", db.Name)
2295+
defer testutil.CleanupDatabase(t, db.Client, dbName)
2296+
2297+
ctx := context.Background()
2298+
2299+
gc := gomongo.NewClient(db.Client)
2300+
2301+
for _, stmt := range []string{
2302+
"db.users.find().count(true)",
2303+
"db.users.find().count(false)",
2304+
"db.users.find().itcount(true)",
2305+
"db.users.find().size(1)",
2306+
} {
2307+
_, err := gc.Execute(ctx, dbName, stmt)
2308+
require.Error(t, err, "expected %q to be rejected", stmt)
2309+
require.Contains(t, err.Error(), "takes no arguments")
2310+
}
21352311
})
21362312
}
21372313

internal/translator/translate.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,19 @@ func translateCollectionStatement(op *Operation, stmt *ast.CollectionStatement)
8989
if err := extractEstimatedDocumentCountArgs(op, stmt.Args); err != nil {
9090
return nil, err
9191
}
92+
case "count":
93+
// db.coll.count() is deprecated in mongosh in favor of countDocuments /
94+
// estimatedDocumentCount. We honor both halves of that recommendation:
95+
// zero-arg count() routes to estimatedDocumentCount (preserves the fast
96+
// metadata path), and any-arg count() routes to countDocuments.
97+
if len(stmt.Args) == 0 {
98+
op.OpType = types.OpEstimatedDocumentCount
99+
} else {
100+
op.OpType = types.OpCountDocuments
101+
if err := extractCountDocumentsArgs(op, stmt.Args); err != nil {
102+
return nil, err
103+
}
104+
}
92105
case "distinct":
93106
op.OpType = types.OpDistinct
94107
if err := extractDistinctArgs(op, stmt.Args); err != nil {
@@ -239,6 +252,24 @@ func translateCursorMethod(op *Operation, cm ast.CursorMethod) error {
239252
return extractMin(op, cm.Args)
240253
case "pretty":
241254
return nil // no-op
255+
case "count", "itcount", "size":
256+
// mongosh's cursor.count() never iterates the cursor; it issues a
257+
// separate count command server-side. We mirror that by retargeting
258+
// the operation to CountDocuments with the accumulated
259+
// filter+skip+limit+hint. Aggregate cursors also expose itcount() but
260+
// require pipeline rewriting ($count stage); not yet supported.
261+
if len(cm.Args) > 0 {
262+
// mongosh historically accepted cursor.count(applySkipLimit), but
263+
// the boolean is a no-op in modern drivers (skip/limit always
264+
// apply). Reject rather than silently drop, since "no-op" is
265+
// hard to debug.
266+
return fmt.Errorf("%s() takes no arguments", cm.Method)
267+
}
268+
if op.OpType != types.OpFind {
269+
return &UnsupportedOperationError{Operation: cm.Method + "()"}
270+
}
271+
op.OpType = types.OpCountDocuments
272+
return nil
242273
default:
243274
return &UnsupportedOperationError{Operation: cm.Method + "()"}
244275
}

0 commit comments

Comments
 (0)