From ca46570f1597dc910bfa87af463fb8f420ed157c Mon Sep 17 00:00:00 2001
From: Andreas Braun <git@alcaeus.org>
Date: Sat, 19 Aug 2017 21:54:19 +0200
Subject: [PATCH] Apply primary readPreference to certain commands

---
 lib/Mongo/MongoCommandCursor.php              |  49 ++++++-
 .../Mongo/MongoCommandCursorTest.php          | 137 ++++++++++++++++++
 2 files changed, 185 insertions(+), 1 deletion(-)

diff --git a/lib/Mongo/MongoCommandCursor.php b/lib/Mongo/MongoCommandCursor.php
index bf3afe42..a5054ac7 100644
--- a/lib/Mongo/MongoCommandCursor.php
+++ b/lib/Mongo/MongoCommandCursor.php
@@ -64,7 +64,19 @@ protected function ensureCursor()
                 }
             }
 
-            $this->cursor = $this->db->command($convertedCommand, $this->getOptions());
+            $originalReadPreference = null;
+            if (!$this->supportsReadPreference()) {
+                $originalReadPreference = $this->readPreference;
+                $this->setReadPreference(\MongoClient::RP_PRIMARY);
+            }
+
+            try {
+                $this->cursor = $this->db->command($convertedCommand, $this->getOptions());
+            } finally {
+                if ($originalReadPreference) {
+                    $this->readPreference = $originalReadPreference;
+                }
+            }
         }
 
         return $this->cursor;
@@ -112,4 +124,39 @@ public function __sleep()
     {
         return ['command'] + parent::__sleep();
     }
+
+    /**
+     * @see https://github.com/mongodb/mongo-php-driver-legacy/blob/1.6.14/db.c#L51
+     * @return bool
+     */
+    private function supportsReadPreference()
+    {
+        if ($this->command === []) {
+            return false;
+        }
+
+        $firstKey = array_keys($this->command)[0];
+        switch ($firstKey) {
+            case 'count':
+            case 'group':
+            case 'dbStats':
+            case 'geoNear':
+            case 'geoWalk':
+            case 'distinct':
+            case 'aggregate':
+            case 'collStats':
+            case 'geoSearch':
+            case 'parallelCollectionScan':
+                return true;
+
+            case 'mapreduce':
+            case 'mapReduce':
+                return (isset($this->command['out']) &&
+                    is_array($this->command['out']) &&
+                    array_key_exists('inline', $this->command['out']));
+
+            default:
+                return false;
+        }
+    }
 }
diff --git a/tests/Alcaeus/MongoDbAdapter/Mongo/MongoCommandCursorTest.php b/tests/Alcaeus/MongoDbAdapter/Mongo/MongoCommandCursorTest.php
index dd878f5c..e5a8d5c3 100644
--- a/tests/Alcaeus/MongoDbAdapter/Mongo/MongoCommandCursorTest.php
+++ b/tests/Alcaeus/MongoDbAdapter/Mongo/MongoCommandCursorTest.php
@@ -2,6 +2,7 @@
 
 namespace Alcaeus\MongoDbAdapter\Tests\Mongo;
 
+use MongoDB\Database;
 use MongoDB\Driver\ReadPreference;
 use Alcaeus\MongoDbAdapter\Tests\TestCase;
 
@@ -65,4 +66,140 @@ public function testInfo()
             $i++;
         }
     }
+
+    /**
+     * @dataProvider dataCommandAppliesCorrectReadPreference
+     */
+    public function testCommandAppliesCorrectReadPreference($command, $expectedReadPreference)
+    {
+        $this->skipTestIf(extension_loaded('mongo'));
+
+        $checkReadPreference = function ($other) use ($expectedReadPreference) {
+            if (!is_array($other)) {
+                return false;
+            }
+
+            if (!array_key_exists('readPreference', $other)) {
+                return false;
+            }
+
+            if (!$other['readPreference'] instanceof ReadPreference) {
+                return false;
+            }
+
+            return $other['readPreference']->getMode() === $expectedReadPreference;
+        };
+
+        $databaseMock = $this->createMock(Database::class);
+        $databaseMock
+            ->expects($this->once())
+            ->method('command')
+            ->with($this->anything(), $this->callback($checkReadPreference))
+            ->will($this->returnValue(new \ArrayIterator()));
+
+        $cursor = new \MongoCommandCursor($this->getClient(), (string) $this->getDatabase(), $command);
+        $reflection = new \ReflectionProperty($cursor, 'db');
+        $reflection->setAccessible(true);
+        $reflection->setValue($cursor, $databaseMock);
+        $cursor->setReadPreference(\MongoClient::RP_SECONDARY);
+
+        iterator_to_array($cursor);
+
+        self::assertSame(\MongoClient::RP_SECONDARY, $cursor->getReadPreference()['type']);
+    }
+
+    public function dataCommandAppliesCorrectReadPreference()
+    {
+        return [
+            'findAndUpdate' => [
+                [
+                    'findandmodify' => (string) $this->getCollection(),
+                    'query' => [],
+                    'update' => ['$inc' => ['field' => 1]],
+                ],
+                ReadPreference::RP_PRIMARY,
+            ],
+            'findAndRemove' => [
+                [
+                    'findandremove' => (string) $this->getCollection(),
+                    'query' => [],
+                ],
+                ReadPreference::RP_PRIMARY,
+            ],
+            'mapReduceWithOut' => [
+                [
+                    'mapReduce' => (string) $this->getCollection(),
+                    'out' => 'sample',
+                ],
+                ReadPreference::RP_PRIMARY,
+            ],
+            'mapReduceWithOutInline' => [
+                [
+                    'mapReduce' => (string) $this->getCollection(),
+                    'out' => ['inline' => true],
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'count' => [
+                [
+                    'count' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'group' => [
+                [
+                    'group' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'dbStats' => [
+                [
+                    'dbStats' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'geoNear' => [
+                [
+                    'geoNear' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'geoWalk' => [
+                [
+                    'geoWalk' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'distinct' => [
+                [
+                    'distinct' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'aggregate' => [
+                [
+                    'aggregate' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'collStats' => [
+                [
+                    'collStats' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'geoSearch' => [
+                [
+                    'geoSearch' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+            'parallelCollectionScan' => [
+                [
+                    'parallelCollectionScan' => (string) $this->getCollection(),
+                ],
+                ReadPreference::RP_SECONDARY,
+            ],
+        ];
+    }
 }