PYCBC-1594: Replica reads from preferred server group

Motivation
==========
Operations do not currently take into account server groups. For
reading from replica servers we need to provide a way for users to
be able to read from specific groups, which can correspond to
availablity zones. Doing so helps users to reduce costs of sending
requests across availability zones.

Changes
=======
* Add a preferred_server_group option to ClusterOptions and propagate it to the C++ core
* Add a read_preference option in replica read operations and propagate it to the C++ core

Results
=======
All tests pass

Change-Id: Ib98a92d9cf141e4e6dbbbdeb9e43d17aa5811d6e
Reviewed-on: https://review.couchbase.org/c/couchbase-python-client/+/221585
Reviewed-by: Jared Casey <jared.casey@couchbase.com>
Tested-by: Build Bot <build@couchbase.com>

Author: DemetrisChr

acouchbase/tests/collection_t.py

@@ -32,10 +32,13 @@
                                   InvalidArgumentException,
                                   PathNotFoundException,
                                   TemporaryFailException)
-from couchbase.options import (GetOptions,
+from couchbase.options import (GetAllReplicasOptions,
+                               GetAnyReplicaOptions,
+                               GetOptions,
                                InsertOptions,
                                ReplaceOptions,
                                UpsertOptions)
+from couchbase.replica_reads import ReadPreference
 from couchbase.result import (ExistsResult,
                               GetReplicaResult,
                               GetResult,
@@ -94,6 +97,10 @@ def check_xattr_supported(self, cb_env):
     def check_not_locked_supported(self, cb_env):
         cb_env.check_if_feature_supported('kv_not_locked')
 
+    @pytest.fixture(scope='class')
+    def check_server_groups_supported(self, cb_env):
+        cb_env.check_if_feature_supported('server_groups')
+
     @pytest_asyncio.fixture(name="new_kvp")
     async def new_key_and_value_with_reset(self, cb_env) -> KVPair:
         key, value = await cb_env.get_new_key_value()
@@ -626,6 +633,17 @@ async def test_get_any_replica_fail(self, cb_env):
         with pytest.raises(DocumentUnretrievableException):
             await cb_env.collection.get_any_replica('not-a-key')
 
+    @pytest.mark.usefixtures("check_multi_node")
+    @pytest.mark.usefixtures("check_replicas")
+    @pytest.mark.usefixtures("check_server_groups_supported")
+    @pytest.mark.asyncio
+    async def test_get_any_replica_read_preference(self, cb_env, default_kvp):
+        # No preferred server group was specified in the cluster options so this should raise
+        # DocumentUnretrievableException
+        with pytest.raises(DocumentUnretrievableException):
+            await cb_env.collection.get_any_replica(
+                default_kvp.key, GetAnyReplicaOptions(read_preference=ReadPreference.SELECTED_SERVER_GROUP))
+
     @pytest.mark.usefixtures("check_multi_node")
     @pytest.mark.usefixtures("check_replicas")
     @pytest.mark.asyncio
@@ -667,6 +685,17 @@ async def test_get_all_replicas_results(self, cb_env, default_kvp):
         assert active_cnt == 1
         assert replica_cnt >= active_cnt
 
+    @pytest.mark.usefixtures("check_multi_node")
+    @pytest.mark.usefixtures("check_replicas")
+    @pytest.mark.usefixtures("check_server_groups_supported")
+    @pytest.mark.asyncio
+    async def test_get_all_replicas_read_preference(self, cb_env, default_kvp):
+        # No preferred server group was specified in the cluster options so this should raise
+        # DocumentUnretrievableException
+        with pytest.raises(DocumentUnretrievableException):
+            await cb_env.collection.get_all_replicas(
+                default_kvp.key, GetAllReplicasOptions(read_preference=ReadPreference.SELECTED_SERVER_GROUP))
+
     # @TODO(jc): - should an expiry of -1 raise an InvalidArgumentException?
     @pytest.mark.usefixtures("check_xattr_supported")
     @pytest.mark.asyncio

acouchbase/tests/subdoc_t.py

@@ -32,6 +32,7 @@
                                LookupInAnyReplicaOptions,
                                LookupInOptions,
                                MutateInOptions)
+from couchbase.replica_reads import ReadPreference
 from couchbase.result import (GetResult,
                               LookupInReplicaResult,
                               LookupInResult,
@@ -72,6 +73,7 @@ class SubDocumentTestSuite:
         'test_lookup_in_all_replicas_get_full',
         'test_lookup_in_all_replicas_multiple_specs',
         'test_lookup_in_all_replicas_with_timeout',
+        'test_lookup_in_all_replicas_read_preference',
         'test_lookup_in_any_replica_bad_key',
         'test_lookup_in_any_replica_exists',
         'test_lookup_in_any_replica_exists_bad_path',
@@ -80,6 +82,7 @@ class SubDocumentTestSuite:
         'test_lookup_in_any_replica_get_full',
         'test_lookup_in_any_replica_multiple_specs',
         'test_lookup_in_any_replica_with_timeout',
+        'test_lookup_in_any_replica_read_preference',
         'test_lookup_in_multiple_specs',
         'test_lookup_in_one_path_not_found',
         'test_lookup_in_simple_exists',
@@ -134,6 +137,13 @@ def check_replica_read_supported(self, cb_env):
                                                        cb_env.server_version_short,
                                                        cb_env.mock_server_type)
 
+    @pytest.fixture(scope='class')
+    def check_server_groups_supported(self, cb_env):
+        EnvironmentFeatures.check_if_feature_supported('server_groups',
+                                                       cb_env.server_version_short,
+                                                       cb_env.mock_server_type,
+                                                       cb_env.server_version_patch)
+
     @pytest.mark.asyncio
     @pytest.mark.usefixtures('skip_if_go_caves')
     async def test_array_add_unique(self, cb_env):
@@ -455,6 +465,19 @@ async def test_lookup_in_all_replicas_with_timeout(self, cb_env):
             active_count += not result.is_replica
         assert active_count == 1
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures('check_replica_read_supported')
+    @pytest.mark.usefixtures('check_server_groups_supported')
+    async def test_lookup_in_all_replicas_read_preference(self, cb_env):
+        key, value = cb_env.get_existing_doc_by_type('vehicle')
+        # No preferred server group was specified in the cluster options so this should raise
+        # DocumentUnretrievableException
+        with pytest.raises(DocumentUnretrievableException):
+            await cb_env.collection.lookup_in_all_replicas(
+                key,
+                [SD.get('batch')],
+                LookupInAllReplicasOptions(read_preference=ReadPreference.SELECTED_SERVER_GROUP))
+
     @pytest.mark.asyncio
     @pytest.mark.usefixtures('check_replica_read_supported')
     async def test_lookup_in_any_replica_bad_key(self, cb_env):
@@ -527,6 +550,17 @@ async def test_lookup_in_any_replica_with_timeout(self, cb_env):
         assert result.content_as[str](0) == value['batch']
         assert result.is_replica is not None
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures('check_replica_read_supported')
+    @pytest.mark.usefixtures('check_server_groups_supported')
+    async def test_lookup_in_any_replica_read_preference(self, cb_env):
+        key, value = cb_env.get_existing_doc_by_type('vehicle')
+        # No preferred server group was specified in the cluster options so this should raise
+        # DocumentUnretrievableException
+        with pytest.raises(DocumentUnretrievableException):
+            await cb_env.collection.lookup_in_any_replica(
+                key, [SD.get('batch')], LookupInAnyReplicaOptions(read_preference=ReadPreference.SELECTED_SERVER_GROUP))
+
     @pytest.mark.asyncio
     @pytest.mark.usefixtures("check_xattr_supported")
     async def test_lookup_in_multiple_specs(self, cb_env):

couchbase/logic/options.py

@@ -46,6 +46,7 @@
     from couchbase.metrics import CouchbaseMeter
     from couchbase.mutation_state import MutationState
     from couchbase.n1ql import QueryProfile, QueryScanConsistency
+    from couchbase.replica_reads import ReadPreference
     from couchbase.search import (Facet,
                                   HighlightStyle,
                                   SearchScanConsistency,
@@ -119,6 +120,7 @@ def get_valid_args(
     'project': lambda x: x,
     'delta': lambda x: x,
     'initial': lambda x: x,
+    'read_preference': lambda x: x.value,
     'per_key_options': lambda x: x,
     'return_exceptions': validate_bool
 }
@@ -556,6 +558,7 @@ class ClusterOptionsBase(dict):
         "dns_nameserver": {"dns_nameserver": validate_str},
         "dns_port": {"dns_port": validate_int},
         "dump_configuration": {"dump_configuration": validate_bool},
+        "preferred_server_group": {"preferred_server_group": validate_str},
     }
 
     @overload
@@ -596,7 +599,8 @@ def __init__(
         dns_nameserver=None,  # type: Optional[str]
         dns_port=None,  # type: Optional[int]
         disable_mozilla_ca_certificates=None,  # type: Optional[bool]
-        dump_configuration=None,  # type: Optional[bool]
+        dump_configuration=None,  # type: Optional[bool],
+        preferred_server_group=None,  # type: Optional[str]
     ):
         """ClusterOptions instance."""
 
@@ -913,7 +917,8 @@ class GetAllReplicasOptionsBase(OptionsTimeoutBase):
     @overload
     def __init__(self,
                  timeout=None,  # type: Optional[timedelta]
-                 transcoder=None  # type: Optional[Transcoder]
+                 transcoder=None,  # type: Optional[Transcoder]
+                 read_preference=None,  # type: Optional[ReadPreference]
                  ):
         pass
 
@@ -952,7 +957,8 @@ class GetAnyReplicaOptionsBase(OptionsTimeoutBase):
     @overload
     def __init__(self,
                  timeout=None,  # type: Optional[timedelta]
-                 transcoder=None  # type: Optional[Transcoder]
+                 transcoder=None,  # type: Optional[Transcoder]
+                 read_preference=None,  # type: Optional[ReadPreference]
                  ):
         pass
 
@@ -995,6 +1001,7 @@ def __init__(self,
                  timeout=None,  # type: Optional[timedelta]
                  span=None,  # type: Optional[Any]
                  serializer=None,  # type: Optional[Serializer]
+                 read_preference=None,  # type: Optional[ReadPreference]
                  ) -> None:
         pass
 
@@ -1009,6 +1016,7 @@ def __init__(self,
                  timeout=None,  # type: Optional[timedelta]
                  span=None,  # type: Optional[Any]
                  serializer=None,  # type: Optional[Serializer]
+                 read_preference=None,  # type: Optional[ReadPreference]
                  ) -> None:
         pass
 

couchbase/options.py

@@ -96,6 +96,7 @@
     from couchbase.collection import Collection
     from couchbase.durability import DurabilityType, ServerDurability
     from couchbase.n1ql import QueryScanConsistency
+    from couchbase.replica_reads import ReadPreference
     from couchbase.transactions import TransactionKeyspace
     from couchbase.transcoder import Transcoder
 
@@ -362,6 +363,7 @@ class ClusterOptions(ClusterOptionsBase):
         dns_nameserver (str, optional):  **VOLATILE** This API is subject to change at any time. Set to configure custom DNS nameserver. Defaults to None.
         dns_port (int, optional):  **VOLATILE** This API is subject to change at any time. Set to configure custom DNS port. Defaults to None.
         dump_configuration (bool, optional): Set to True to dump every new configuration when TRACE level logging. Defaults to False (disabled).
+        preferred_server_group (str, optional): Specifies the preferred server group to be used for replica reads with 'selected server group' read preference.
     """  # noqa: E501
 
     def apply_profile(self,
@@ -507,6 +509,8 @@ class GetAllReplicasOptions(GetAllReplicasOptionsBase):
             key-value operation timeout.
         transcoder (:class:`~.transcoder.Transcoder`, optional): Specifies an explicit transcoder
             to use for this specific operation. Defaults to :class:`~.transcoder.JsonTranscoder`.
+        read_preference(:class:`~couchbase.replica_reads.ReadPreference`, optional): Specifies how the replica nodes
+            will be selected. Defaults to no preference.
     """
 
 
@@ -552,6 +556,8 @@ class GetAnyReplicaOptions(GetAnyReplicaOptionsBase):
             key-value operation timeout.
         transcoder (:class:`~.transcoder.Transcoder`, optional): Specifies an explicit transcoder
             to use for this specific operation. Defaults to :class:`~.transcoder.JsonTranscoder`.
+        read_preference(:class:`~couchbase.replica_reads.ReadPreference`, optional): Specifies how the replica nodes
+            will be selected. Defaults to no preference.
     """
 
 
@@ -706,6 +712,8 @@ class LookupInAnyReplicaOptions(LookupInAnyReplicaOptionsBase):
     Args:
         timeout (timedelta, optional): The timeout for this operation. Defaults to global
             subdocument operation timeout.
+        read_preference(:class:`~couchbase.replica_reads.ReadPreference`, optional): Specifies how the replica nodes
+            will be selected. Defaults to no preference.
     """
 
 
@@ -719,6 +727,8 @@ class LookupInAllReplicasOptions(LookupInAllReplicasOptionsBase):
     Args:
         timeout (timedelta, optional): The timeout for this operation. Defaults to global
             subdocument operation timeout.
+        read_preference(:class:`~couchbase.replica_reads.ReadPreference`, optional): Specifies how the replica nodes
+            will be selected. Defaults to no preference.
     """
 
 
@@ -834,6 +844,8 @@ class GetAllReplicasMultiOptions(dict):
             key-value operation timeout.
         transcoder (:class:`~couchbase.transcoder.Transcoder`, optional): Specifies an explicit transcoder
             to use for this specific operation. Defaults to :class:`~.transcoder.JsonTranscoder`.
+        read_preference(:class:`~couchbase.replica_reads.ReadPreference`, optional): Specifies how the replica nodes
+            will be selected. Defaults to no preference.
         per_key_options (Dict[str, :class:`.GetAllReplicasOptions`], optional): Specify
             :class:`.GetAllReplicasOptions` per key.
         return_exceptions(bool, optional): If False, raise an Exception when encountered.  If True return the
@@ -842,9 +854,10 @@ class GetAllReplicasMultiOptions(dict):
     @overload
     def __init__(
         self,
-        transcoder=None,  # type: Transcoder
-        per_key_options=None,       # type: Dict[str, GetAllReplicasOptions]
-        return_exceptions=None      # type: Optional[bool]
+        transcoder=None,        # type: Optional[Transcoder]
+        read_preference=None,   # type: Optional[ReadPreference]
+        per_key_options=None,   # type: Dict[str, GetAllReplicasOptions]
+        return_exceptions=None  # type: Optional[bool]
     ):
         pass
 
@@ -854,7 +867,7 @@ def __init__(self, **kwargs):
 
     @classmethod
     def get_valid_keys(cls):
-        return ['timeout', 'transcoder', 'per_key_options', 'return_exceptions']
+        return ['timeout', 'transcoder', 'read_preference', 'per_key_options', 'return_exceptions']
 
 
 class GetAnyReplicaMultiOptions(dict):
@@ -868,6 +881,8 @@ class GetAnyReplicaMultiOptions(dict):
             key-value operation timeout.
         transcoder (:class:`~couchbase.transcoder.Transcoder`, optional): Specifies an explicit transcoder
             to use for this specific operation. Defaults to :class:`~.transcoder.JsonTranscoder`.
+        read_preference(:class:`~couchbase.replica_reads.ReadPreference`, optional): Specifies how the replica nodes
+            will be selected. Defaults to no preference.
         per_key_options (Dict[str, :class:`.GetAnyReplicaOptions`], optional): Specify
             :class:`.GetAnyReplicaOptions` per key.
         return_exceptions(bool, optional): If False, raise an Exception when encountered.  If True return the
@@ -876,9 +891,10 @@ class GetAnyReplicaMultiOptions(dict):
     @overload
     def __init__(
         self,
-        transcoder=None,  # type: Transcoder
-        per_key_options=None,       # type: Dict[str, GetAnyReplicaOptions]
-        return_exceptions=None      # type: Optional[bool]
+        transcoder=None,        # type: Optional[Transcoder]
+        read_preference=None,   # type: Optional[ReadPreference]
+        per_key_options=None,   # type: Dict[str, GetAnyReplicaOptions]
+        return_exceptions=None  # type: Optional[bool]
     ):
         pass
 
@@ -888,7 +904,7 @@ def __init__(self, **kwargs):
 
     @classmethod
     def get_valid_keys(cls):
-        return ['timeout', 'transcoder', 'per_key_options', 'return_exceptions']
+        return ['timeout', 'transcoder', 'read_preference', 'per_key_options', 'return_exceptions']
 
 
 class GetMultiOptions(dict):
@@ -1820,7 +1836,8 @@ def arg_mapping(self):
             "report_id": {"report_id": lambda x: str(x)},
             "batch_byte_limit": {"batch_byte_limit": validate_int},
             "batch_item_limit": {"batch_item_limit": validate_int},
-            "concurrency": {"concurrency": validate_int}
+            "concurrency": {"concurrency": validate_int},
+            "read_preference": {"read_preference": lambda r: r.value}
         }
 
 

couchbase/replica_reads.py

@@ -0,0 +1,6 @@
+from enum import IntEnum
+
+
+class ReadPreference(IntEnum):
+    NO_PREFERENCE = 0
+    SELECTED_SERVER_GROUP = 1