Skip to content

Commit 986316c

Browse files
authored
CDRIVER-4316 Allow both int32 and int64 for batchSize in cursor opts (#1300)
1 parent a41683b commit 986316c

File tree

3 files changed

+318
-3
lines changed

3 files changed

+318
-3
lines changed

src/libmongoc/src/mongoc/mongoc-cursor.c

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1449,9 +1449,32 @@ void
14491449
mongoc_cursor_set_batch_size (mongoc_cursor_t *cursor, uint32_t batch_size)
14501450
{
14511451
BSON_ASSERT (cursor);
1452-
1453-
_mongoc_cursor_set_opt_int64 (
1454-
cursor, MONGOC_CURSOR_BATCH_SIZE, (int64_t) batch_size);
1452+
bson_iter_t iter;
1453+
if (!bson_iter_init_find (&iter, &cursor->opts, MONGOC_CURSOR_BATCH_SIZE)) {
1454+
bson_append_int64 (&cursor->opts,
1455+
MONGOC_CURSOR_BATCH_SIZE,
1456+
MONGOC_CURSOR_BATCH_SIZE_LEN,
1457+
batch_size);
1458+
} else if (BSON_ITER_HOLDS_INT64 (&iter)) {
1459+
bson_iter_overwrite_int64 (&iter, (int64_t) batch_size);
1460+
} else if (BSON_ITER_HOLDS_INT32 (&iter)) {
1461+
if (!bson_in_range_int32_t_unsigned (batch_size)) {
1462+
MONGOC_WARNING ("unable to overwrite stored int32 batchSize with "
1463+
"out-of-range value %" PRIu32,
1464+
batch_size);
1465+
return;
1466+
}
1467+
bson_iter_overwrite_int32 (&iter, (int32_t) batch_size);
1468+
} else if (BSON_ITER_HOLDS_DOUBLE (&iter)) {
1469+
bson_iter_overwrite_double (&iter, (double) batch_size);
1470+
} else if (BSON_ITER_HOLDS_DECIMAL128 (&iter)) {
1471+
bson_decimal128_t val;
1472+
val.high = 0x3040000000000000;
1473+
val.low = (uint64_t) batch_size;
1474+
bson_iter_overwrite_decimal128 (&iter, &val);
1475+
} else {
1476+
MONGOC_WARNING ("unable to overwrite non-numeric stored batchSize");
1477+
}
14551478
}
14561479

14571480

src/libmongoc/tests/test-mongoc-collection-find-with-opts.c

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -840,6 +840,88 @@ test_server_id_option (void)
840840
mongoc_client_destroy (client);
841841
}
842842

843+
static void
844+
test_find_batchSize (void)
845+
{
846+
mongoc_client_t *client;
847+
mongoc_collection_t *collection;
848+
bson_error_t error;
849+
mongoc_cursor_t *cursor;
850+
851+
client = test_framework_new_default_client ();
852+
collection = mongoc_client_get_collection (client, "db", "collection");
853+
854+
// Test a cursor with an int32 batchSize.
855+
{
856+
cursor = mongoc_collection_find_with_opts (
857+
collection,
858+
tmp_bson ("{}"),
859+
tmp_bson ("{'batchSize': { '$numberInt': '1' }}"),
860+
NULL);
861+
862+
ASSERT_OR_PRINT (!mongoc_cursor_error (cursor, &error), error);
863+
864+
// Exhaust the cursor.
865+
{
866+
// Note: Cursors are lazy. The `find` command is sent on the first call
867+
// to `mongoc_cursor_next`.
868+
const bson_t *got;
869+
while (mongoc_cursor_next (cursor, &got))
870+
;
871+
}
872+
873+
ASSERT_OR_PRINT (!mongoc_cursor_error (cursor, &error), error);
874+
mongoc_cursor_destroy (cursor);
875+
}
876+
877+
// Test a cursor with an int64 batchSize.
878+
{
879+
cursor = mongoc_collection_find_with_opts (
880+
collection,
881+
tmp_bson ("{}"),
882+
tmp_bson ("{'batchSize': { '$numberLong': '1' }}"),
883+
NULL);
884+
885+
ASSERT_OR_PRINT (!mongoc_cursor_error (cursor, &error), error);
886+
887+
// Exhaust the cursor.
888+
{
889+
// Note: Cursors are lazy. The `find` command is sent on the first call
890+
// to `mongoc_cursor_next`.
891+
const bson_t *got;
892+
while (mongoc_cursor_next (cursor, &got))
893+
;
894+
}
895+
896+
ASSERT_OR_PRINT (!mongoc_cursor_error (cursor, &error), error);
897+
mongoc_cursor_destroy (cursor);
898+
}
899+
900+
// Test a cursor with a string batchSize.
901+
{
902+
cursor = mongoc_collection_find_with_opts (
903+
collection, tmp_bson ("{}"), tmp_bson ("{'batchSize': 'foo'}"), NULL);
904+
905+
ASSERT_OR_PRINT (!mongoc_cursor_error (cursor, &error), error);
906+
907+
// Attempt to exhaust the cursor.
908+
{
909+
// Note: Cursors are lazy. The `find` command is sent on the first call
910+
// to `mongoc_cursor_next`.
911+
const bson_t *got;
912+
while (mongoc_cursor_next (cursor, &got))
913+
;
914+
}
915+
916+
// Expect an error from the server.
917+
ASSERT (mongoc_cursor_error (cursor, &error));
918+
mongoc_cursor_destroy (cursor);
919+
}
920+
921+
mongoc_collection_destroy (collection);
922+
mongoc_client_destroy (client);
923+
}
924+
843925
void
844926
test_collection_find_with_opts_install (TestSuite *suite)
845927
{
@@ -920,4 +1002,5 @@ test_collection_find_with_opts_install (TestSuite *suite)
9201002
TestSuite_AddLive (suite,
9211003
"/Collection/find_with_opts/server_id/option",
9221004
test_server_id_option);
1005+
TestSuite_AddLive (suite, "/Collection/find/batchSize", test_find_batchSize);
9231006
}

src/libmongoc/tests/test-mongoc-cursor.c

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2352,6 +2352,200 @@ test_find_error_is_alive (void)
23522352
mongoc_client_destroy (client);
23532353
}
23542354

2355+
typedef struct _started_event_t {
2356+
char *command_name;
2357+
bson_t *command;
2358+
} started_event_t;
2359+
2360+
static void
2361+
command_started (const mongoc_apm_command_started_t *event)
2362+
{
2363+
mongoc_array_t *events =
2364+
(mongoc_array_t *) mongoc_apm_command_started_get_context (event);
2365+
started_event_t *started_event = bson_malloc0 (sizeof (started_event_t));
2366+
2367+
started_event->command =
2368+
bson_copy (mongoc_apm_command_started_get_command (event));
2369+
started_event->command_name =
2370+
bson_strdup (mongoc_apm_command_started_get_command_name (event));
2371+
_mongoc_array_append_val (events, started_event);
2372+
}
2373+
2374+
static void
2375+
clear_started_events (mongoc_array_t *events)
2376+
{
2377+
for (size_t i = 0; i < events->len; i++) {
2378+
started_event_t *started_event =
2379+
_mongoc_array_index (events, started_event_t *, i);
2380+
bson_destroy (started_event->command);
2381+
bson_free (started_event->command_name);
2382+
bson_free (started_event);
2383+
}
2384+
_mongoc_array_clear (events);
2385+
}
2386+
2387+
void
2388+
numeric_iter_eq (bson_iter_t *iter, int64_t val)
2389+
{
2390+
ASSERT_CMPINT64 (bson_iter_as_int64 (iter), ==, val);
2391+
}
2392+
2393+
void
2394+
decimal128_iter_eq (bson_iter_t *iter, int64_t val)
2395+
{
2396+
bson_decimal128_t d;
2397+
bson_iter_decimal128 (iter, &d);
2398+
ASSERT_CMPUINT64 (d.high, ==, 0x3040000000000000);
2399+
ASSERT_CMPINT64 (d.low, ==, val);
2400+
}
2401+
2402+
void
2403+
test_cursor_batchsize_override (bson_t *findopts,
2404+
void (*assert_eq) (bson_iter_t *, int64_t))
2405+
{
2406+
mongoc_client_t *client;
2407+
mongoc_apm_callbacks_t *cbs;
2408+
mongoc_collection_t *coll;
2409+
bson_error_t error;
2410+
mongoc_array_t started_events;
2411+
2412+
client = test_framework_new_default_client ();
2413+
cbs = mongoc_apm_callbacks_new ();
2414+
_mongoc_array_init (&started_events, sizeof (started_event_t *));
2415+
mongoc_apm_set_command_started_cb (cbs, command_started);
2416+
coll = mongoc_client_get_collection (client, "db", "coll");
2417+
2418+
/* Drop and insert two documents into the collection */
2419+
{
2420+
bson_t *to_insert = BCON_NEW ("x", "y");
2421+
2422+
// Ignore "ns not found" error on drop.
2423+
mongoc_collection_drop (coll, NULL);
2424+
ASSERT_OR_PRINT (
2425+
mongoc_collection_insert_one (
2426+
coll, to_insert, NULL /* opts */, NULL /* reply */, &error),
2427+
error);
2428+
ASSERT_OR_PRINT (
2429+
mongoc_collection_insert_one (
2430+
coll, to_insert, NULL /* opts */, NULL /* reply */, &error),
2431+
error);
2432+
bson_destroy (to_insert);
2433+
}
2434+
2435+
mongoc_client_set_apm_callbacks (client, cbs, &started_events);
2436+
2437+
/* Create a cursor and iterate once. */
2438+
{
2439+
const bson_t *got;
2440+
bson_t *filter = bson_new ();
2441+
mongoc_cursor_t *cursor = mongoc_collection_find_with_opts (
2442+
coll, filter, findopts, NULL /* read_prefs */);
2443+
/* Attempt to overwrite the 'batchSize' with 2. */
2444+
mongoc_cursor_set_batch_size (cursor, 2);
2445+
/* Assert no command started events. The cursor does not send 'find' until
2446+
* the first call to mongoc_cursor_next. */
2447+
ASSERT_CMPSIZE_T (started_events.len, ==, 0);
2448+
/* Iterate once. */
2449+
ASSERT (mongoc_cursor_next (cursor, &got));
2450+
2451+
mongoc_cursor_destroy (cursor);
2452+
bson_destroy (findopts);
2453+
bson_destroy (filter);
2454+
}
2455+
2456+
/* Check events. */
2457+
{
2458+
started_event_t *started_event;
2459+
bson_iter_t iter;
2460+
/* Expect first event is find. */
2461+
started_event =
2462+
_mongoc_array_index (&started_events, started_event_t *, 0);
2463+
ASSERT_CMPSTR (started_event->command_name, "find");
2464+
/* Expect the batchSize sent to be 2. */
2465+
ASSERT (bson_iter_init_find (&iter, started_event->command, "batchSize"));
2466+
assert_eq (&iter, 2);
2467+
}
2468+
2469+
mongoc_collection_destroy (coll);
2470+
mongoc_apm_callbacks_destroy (cbs);
2471+
mongoc_client_destroy (client);
2472+
2473+
clear_started_events (&started_events);
2474+
_mongoc_array_destroy (&started_events);
2475+
}
2476+
2477+
/* Test that mongoc_cursor_set_batch_size overrides a previously set int32
2478+
* batchSize. */
2479+
void
2480+
test_cursor_batchsize_override_int32 (void)
2481+
{
2482+
bson_t *findopts = BCON_NEW ("batchSize", BCON_INT32 (1));
2483+
test_cursor_batchsize_override (findopts, numeric_iter_eq);
2484+
}
2485+
2486+
/* Test that mongoc_cursor_set_batch_size overrides a previously set int64
2487+
* batchSize. */
2488+
void
2489+
test_cursor_batchsize_override_int64 (void)
2490+
{
2491+
bson_t *findopts = BCON_NEW ("batchSize", BCON_INT64 (1));
2492+
test_cursor_batchsize_override (findopts, numeric_iter_eq);
2493+
}
2494+
2495+
/* Test that mongoc_cursor_set_batch_size overrides a previously set double
2496+
* batchSize. */
2497+
void
2498+
test_cursor_batchsize_override_double (void)
2499+
{
2500+
bson_t *findopts = BCON_NEW ("batchSize", BCON_DOUBLE (1.0));
2501+
test_cursor_batchsize_override (findopts, numeric_iter_eq);
2502+
}
2503+
2504+
/* Test that mongoc_cursor_set_batch_size overrides a previously set decimal128
2505+
* batchSize. */
2506+
void
2507+
test_cursor_batchsize_override_decimal128 (void)
2508+
{
2509+
bson_decimal128_t start_val;
2510+
bson_decimal128_from_string ("1", &start_val);
2511+
bson_t *findopts = BCON_NEW ("batchSize", BCON_DECIMAL128 (&start_val));
2512+
test_cursor_batchsize_override (findopts, decimal128_iter_eq);
2513+
}
2514+
2515+
/* Test that attempting to overwrite an int32 batchSize with an out-of-range
2516+
* value raises a warning */
2517+
void
2518+
test_cursor_batchsize_override_range_warning (void)
2519+
{
2520+
mongoc_client_t *client;
2521+
mongoc_collection_t *coll;
2522+
bson_t *findopts = BCON_NEW ("batchSize", BCON_INT32 (1.0));
2523+
2524+
client = test_framework_new_default_client ();
2525+
coll = mongoc_client_get_collection (client, "db", "coll");
2526+
2527+
/* Create a cursor and attempt to override outside int32 range. */
2528+
{
2529+
bson_t *filter = bson_new ();
2530+
mongoc_cursor_t *cursor = mongoc_collection_find_with_opts (
2531+
coll, filter, findopts, NULL /* read_prefs */);
2532+
2533+
capture_logs (true);
2534+
/* Attempt to overwrite the 'batchSize' with uint32_max. */
2535+
mongoc_cursor_set_batch_size (cursor, UINT32_MAX);
2536+
ASSERT_CAPTURED_LOG (
2537+
"mongoc_cursor_set_batch_size",
2538+
MONGOC_LOG_LEVEL_WARNING,
2539+
"unable to overwrite stored int32 batchSize with out-of-range value");
2540+
2541+
mongoc_cursor_destroy (cursor);
2542+
bson_destroy (findopts);
2543+
bson_destroy (filter);
2544+
}
2545+
2546+
mongoc_collection_destroy (coll);
2547+
mongoc_client_destroy (client);
2548+
}
23552549

23562550
void
23572551
test_cursor_install (TestSuite *suite)
@@ -2438,4 +2632,19 @@ test_cursor_install (TestSuite *suite)
24382632
suite, "/Cursor/error_document/command", test_error_document_command);
24392633
TestSuite_AddLive (
24402634
suite, "/Cursor/find_error/is_alive", test_find_error_is_alive);
2635+
TestSuite_AddLive (suite,
2636+
"/Cursor/batchsize_override_int32",
2637+
test_cursor_batchsize_override_int32);
2638+
TestSuite_AddLive (suite,
2639+
"/Cursor/batchsize_override_int64",
2640+
test_cursor_batchsize_override_int64);
2641+
TestSuite_AddLive (suite,
2642+
"/Cursor/batchsize_override_double",
2643+
test_cursor_batchsize_override_double);
2644+
TestSuite_AddLive (suite,
2645+
"/Cursor/batchsize_override_decimal128",
2646+
test_cursor_batchsize_override_decimal128);
2647+
TestSuite_AddLive (suite,
2648+
"/Cursor/batchsize_override_range_warning",
2649+
test_cursor_batchsize_override_range_warning);
24412650
}

0 commit comments

Comments
 (0)