SERVER-59071 Treat '$sample' as unsharded when connecting directly to shards

1 files changed, 23 insertions, 0 deletions

diff --git a/jstests/sharding/sample_direct_connection.js b/jstests/sharding/sample_direct_connection.js
new file mode 100644
index 00000000000..c435490c2f8
--- /dev/null
+++ b/jstests/sharding/sample_direct_connection.js
@@ -0,0 +1,23 @@
+// This test ensures that an aggregation pipeline with a $sample stage can be issued through a
+// direct connection to a shard without failing, and $sample behaves as if we sampled an unsharded
+// collection.
+(function() {
+let st = new ShardingTest({shards: 1});
+
+assert.commandWorked(st.adminCommand({enableSharding: 'test'}));
+assert.commandWorked(st.adminCommand({shardCollection: 'test.sharded', key: {x: 1}}));
+const testDB = st.s.getDB('test');
+
+// We must have >100 samples to attempt the storage optimized sample path.
+for (let x = 0; x < 101; x++) {
+    assert.commandWorked(testDB.foo.insert({x: x}));
+}
+
+const shardDB = st.rs0.getPrimary().getDB('test');
+
+const res = assert.commandWorked(
+    shardDB.runCommand({aggregate: 'foo', pipeline: [{$sample: {size: 3}}], cursor: {}}));
+assert.eq(res.cursor.firstBatch.length, 3);
+
+st.stop();
+})();