fix(mango): covering indexes for partitioned databases

The previous work that introduced the keys-only covering indexes
did not count with the case that database might be partitioned.
And since they use a different format for their own local indexes
and the code does not handle that, it will crash.

When indexes are defined globally for the partitioned databases,
there is no problem because the view row does not include
information about the partition, i.e. it is transparent.

Add the missing support for these scenarios and extend the test
suite to cover them as well.  That latter required some changes to
the base classes in the integration test suite as it apparently
completely misses out on running test cases for partitioned
databases.

4 files changed, 125 insertions, 28 deletions

diff --git a/src/mango/src/mango_cursor_view.erl b/src/mango/src/mango_cursor_view.erl
index e044c56fc..82b4008bb 100644
--- a/src/mango/src/mango_cursor_view.erl
+++ b/src/mango/src/mango_cursor_view.erl
@@ -435,7 +435,12 @@ match_and_extract_doc(Doc, Selector, Fields) ->
     end.
 
 -spec derive_doc_from_index(#idx{}, #view_row{}) -> term().
-derive_doc_from_index(Index, #view_row{id = DocId, key = Keys}) ->
+derive_doc_from_index(Index, #view_row{id = DocId, key = KeyData}) ->
+    Keys =
+        case KeyData of
+            {p, _Partition, KeyValues} -> KeyValues;
+            KeyValues -> KeyValues
+        end,
     Columns = mango_idx:columns(Index),
     lists:foldr(
         fun({Column, Key}, Doc) -> mango_doc:set_field(Doc, Column, Key) end,
@@ -862,6 +867,18 @@ derive_doc_from_index_test() ->
     Doc = {[{<<"_id">>, DocId}, {<<"field2">>, key2}, {<<"field1">>, key1}]},
     ?assertEqual(Doc, derive_doc_from_index(Index, ViewRow)).
 
+derive_doc_from_index_partitioned_test() ->
+    Index =
+        #idx{
+            type = <<"json">>,
+            def = {[{<<"fields">>, {[{<<"field1">>, undefined}, {<<"field2">>, undefined}]}}]}
+        },
+    DocId = doc_id,
+    Keys = [key1, key2],
+    ViewRow = #view_row{id = DocId, key = {p, partition, Keys}},
+    Doc = {[{<<"_id">>, DocId}, {<<"field2">>, key2}, {<<"field1">>, key1}]},
+    ?assertEqual(Doc, derive_doc_from_index(Index, ViewRow)).
+
 composite_indexes_test() ->
     ?assertEqual([], composite_indexes([], [])),
     Index1 =
diff --git a/src/mango/test/22-covering-index-test.py b/src/mango/test/22-covering-index-test.py
index 52a7f3612..176dd2d81 100644
--- a/src/mango/test/22-covering-index-test.py
+++ b/src/mango/test/22-covering-index-test.py
@@ -13,19 +13,7 @@
 import mango
 
 
-class CoveringIndexTests(mango.UserDocsTests):
-    def is_covered(self, selector, fields, index, use_index=None):
-        resp = self.db.find(selector, fields=fields, use_index=use_index, explain=True)
-        self.assertEqual(resp["index"]["type"], "json")
-        self.assertEqual(resp["index"]["name"], index)
-        self.assertEqual(resp["mrargs"]["include_docs"], False)
-        self.assertEqual(resp["covered"], True)
-
-    def is_not_covered(self, selector, fields, use_index=None):
-        resp = self.db.find(selector, fields=fields, use_index=use_index, explain=True)
-        self.assertEqual(resp["mrargs"]["include_docs"], True)
-        self.assertEqual(resp["covered"], False)
-
+class CoveringIndexTests:
     def test_index_covers_query_1field_index_id(self):
         self.is_covered({"age": {"$gte": 32}}, ["_id"], "age")
 
@@ -88,6 +76,45 @@ class CoveringIndexTests(mango.UserDocsTests):
             use_index="twitter",
         )
 
+
+class RegularCoveringIndexTests(mango.UserDocsTests, CoveringIndexTests):
+    @classmethod
+    def setUpClass(klass):
+        super(RegularCoveringIndexTests, klass).setUpClass()
+
+    def is_covered(self, selector, fields, index, use_index=None):
+        resp = self.db.find(selector, fields=fields, use_index=use_index, explain=True)
+        self.assertEqual(resp["index"]["type"], "json")
+        self.assertEqual(resp["index"]["name"], index)
+        self.assertEqual(resp["mrargs"]["include_docs"], False)
+        self.assertEqual(resp["covered"], True)
+
+    def is_not_covered(self, selector, fields, use_index=None):
+        resp = self.db.find(selector, fields=fields, use_index=use_index, explain=True)
+        self.assertEqual(resp["mrargs"]["include_docs"], True)
+        self.assertEqual(resp["covered"], False)
+
+    def test_covering_index_provides_correct_answer_2field_index(self):
+        docs = self.db.find(
+            {"company": {"$exists": True}, "manager": True},
+            sort=[{"company": "asc"}],
+            fields=["company"],
+            use_index="company_and_manager",
+        )
+        expected = [
+            {"company": "Affluex"},
+            {"company": "Globoil"},
+            {"company": "Lyria"},
+            {"company": "Manglo"},
+            {"company": "Myopium"},
+            {"company": "Niquent"},
+            {"company": "Oulu"},
+            {"company": "Prosely"},
+            {"company": "Tasmania"},
+            {"company": "Zialactic"},
+        ]
+        self.assertEqual(docs, expected)
+
     def test_covering_index_provides_correct_answer_id(self):
         docs = self.db.find({"age": {"$gte": 32}}, fields=["_id"])
         expected = [
@@ -107,23 +134,50 @@ class CoveringIndexTests(mango.UserDocsTests):
         ]
         self.assertEqual(docs, expected)
 
+
+class PartitionedCoveringIndexTests(mango.PartitionedUserDocsTests, CoveringIndexTests):
+    @classmethod
+    def setUpClass(klass):
+        super(PartitionedCoveringIndexTests, klass).setUpClass()
+
+    def is_covered(self, selector, fields, index, use_index=None):
+        resp = self.db.find(
+            selector, fields=fields, use_index=use_index, explain=True, partition="0"
+        )
+        self.assertEqual(resp["index"]["type"], "json")
+        self.assertEqual(resp["index"]["name"], index)
+        self.assertEqual(resp["mrargs"]["include_docs"], False)
+        self.assertEqual(resp["covered"], True)
+
+    def is_not_covered(self, selector, fields, use_index=None):
+        resp = self.db.find(
+            selector, fields=fields, use_index=use_index, explain=True, partition="0"
+        )
+        self.assertEqual(resp["mrargs"]["include_docs"], True)
+        self.assertEqual(resp["covered"], False)
+
     def test_covering_index_provides_correct_answer_2field_index(self):
         docs = self.db.find(
             {"company": {"$exists": True}, "manager": True},
             sort=[{"company": "asc"}],
             fields=["company"],
             use_index="company_and_manager",
+            partition="0",
         )
         expected = [
-            {"company": "Affluex"},
-            {"company": "Globoil"},
-            {"company": "Lyria"},
             {"company": "Manglo"},
-            {"company": "Myopium"},
-            {"company": "Niquent"},
             {"company": "Oulu"},
             {"company": "Prosely"},
             {"company": "Tasmania"},
-            {"company": "Zialactic"},
         ]
         self.assertEqual(docs, expected)
+
+    def test_covering_index_provides_correct_answer_id(self):
+        docs = self.db.find({"age": {"$gte": 32}}, fields=["_id"], partition="0")
+        expected = [
+            {"_id": "0:0461444c-e60a-457d-a4bb-b8d811853f21"},
+            {"_id": "0:5b61abc1-a3d3-4092-b9d7-ced90e675536"},
+            {"_id": "0:71562648-6acb-42bc-a182-df6b1f005b09"},
+            {"_id": "0:b31dad3f-ae8b-4f86-8327-dfe8770beb27"},
+        ]
+        self.assertEqual(sorted(docs, key=lambda x: x["_id"]), expected)
diff --git a/src/mango/test/mango.py b/src/mango/test/mango.py
index 81bc082b6..20a40d1b7 100644
--- a/src/mango/test/mango.py
+++ b/src/mango/test/mango.py
@@ -66,10 +66,11 @@ class Database(object):
             parts = [parts]
         return "/".join([self.url] + parts)
 
-    def create(self, q=1, n=1):
+    def create(self, q=1, n=1, partitioned=False):
         r = self.sess.get(self.url)
         if r.status_code == 404:
-            r = self.sess.put(self.url, params={"q": q, "n": n})
+            p = str(partitioned).lower()
+            r = self.sess.put(self.url, params={"q": q, "n": n, "partitioned": p})
             r.raise_for_status()
 
     def delete(self):
@@ -240,6 +241,7 @@ class Database(object):
         return_raw=False,
         update=True,
         executionStats=False,
+        partition=None,
     ):
         body = {
             "selector": selector,
@@ -260,10 +262,14 @@ class Database(object):
         if executionStats == True:
             body["execution_stats"] = True
         body = json.dumps(body)
+        if partition:
+            ppath = "_partition/{}/".format(partition)
+        else:
+            ppath = ""
         if explain:
-            path = self.path("_explain")
+            path = self.path("{}_explain".format(ppath))
         else:
-            path = self.path("_find")
+            path = self.path("{}_find".format(ppath))
         r = self.sess.post(path, data=body)
         r.raise_for_status()
         if explain or return_raw:
@@ -298,9 +304,9 @@ class UsersDbTests(unittest.TestCase):
 
 class DbPerClass(unittest.TestCase):
     @classmethod
-    def setUpClass(klass):
+    def setUpClass(klass, partitioned=False):
         klass.db = Database(random_db_name())
-        klass.db.create(q=1, n=1)
+        klass.db.create(q=1, n=1, partitioned=partitioned)
 
     @classmethod
     def tearDownClass(klass):
@@ -320,6 +326,15 @@ class UserDocsTests(DbPerClass):
         user_docs.setup(klass.db)
 
 
+class PartitionedUserDocsTests(DbPerClass):
+    INDEX_TYPE = "json"
+
+    @classmethod
+    def setUpClass(klass):
+        super(PartitionedUserDocsTests, klass).setUpClass(partitioned=True)
+        user_docs.setup(klass.db, partitioned=True)
+
+
 class UserDocsTestsNoIndexes(DbPerClass):
     INDEX_TYPE = "special"
 
diff --git a/src/mango/test/user_docs.py b/src/mango/test/user_docs.py
index a1b6c6e1a..b2a19eea7 100644
--- a/src/mango/test/user_docs.py
+++ b/src/mango/test/user_docs.py
@@ -58,9 +58,18 @@ def setup_users(db, **kwargs):
     db.save_docs(copy.deepcopy(USERS_DOCS))
 
 
-def setup(db, index_type="view", **kwargs):
+def setup(db, index_type="view", partitioned=False, **kwargs):
     db.recreate()
-    db.save_docs(copy.deepcopy(DOCS))
+    p = str(partitioned).lower()
+    docs = copy.deepcopy(DOCS)
+
+    if partitioned:
+        for index, doc in enumerate(docs):
+            partition = index % PARTITIONS
+            doc["_id"] = "{}:{}".format(partition, doc["_id"])
+
+    db.save_docs(docs, partitioned=p)
+
     if index_type == "view":
         add_view_indexes(db, kwargs)
     elif index_type == "text":
@@ -96,6 +105,8 @@ def add_text_indexes(db, kwargs):
     db.create_text_index(**kwargs)
 
 
+PARTITIONS = 3
+
 DOCS = [
     {
         "_id": "71562648-6acb-42bc-a182-df6b1f005b09",