Allow non-exact matching when filtering INFO keys

So we can ask for "db" and get db0, db1, db2, etc.

Co-authored-by: Oran Agra <oran@redislabs.com>

setup.py

@@ -29,7 +29,7 @@
 
 setup(
     name='redis-info-provider',
-    version='0.12.0',
+    version='0.13.0',
     python_requires='>=2.7',
     package_dir={'': 'src'},
     packages=find_packages('src'),

src/redis_info_provider/info_servicer.py

@@ -46,7 +46,7 @@ class InfoProviderServicer(object):
     """Implements the InfoProvider RPC interface."""
 
     @classmethod
-    def _filter_info(cls, full_info, keys):
+    def _filter_info(cls, full_info, keys, prefix_matching=False):
         # type: (InfoType, AbstractSet[str]) -> InfoType
 
         """
@@ -56,13 +56,23 @@ def _filter_info(cls, full_info, keys):
 
         :param full_info: An info_pb2.Info instance to filter.
         :param keys: A set of strings to match for.
+        :param prefix_matching: If True, we treat each key in `keys` as a prefix of a
+            field in `full_info`
         """
 
+        def is_matching_key(full_key, keys, prefix_matching=False):
+            if not prefix_matching:
+                return full_key in keys
+            else:
+                if full_key.startswith(tuple(keys)):
+                    return True
+                return False
+
         # Optimization: If no keys specified, just return the full info that's already waiting
         if not keys:
             return full_info
 
-        info = {k: full_info[k] for k in keys if k in full_info}
+        info = {k: full_info[k] for k in full_info if is_matching_key(k, keys, prefix_matching=prefix_matching)}
         # always include the metadata
         info['meta'] = full_info['meta']
 
@@ -92,7 +102,7 @@ def _get_shard_with_info(shard_id):
         return shard
 
     @deprecated_alias(key_patterns='keys')
-    def GetInfos(self, shard_ids=(), keys=(), allow_partial=False, max_age=0.0):
+    def GetInfos(self, shard_ids=(), keys=(), allow_partial=False, max_age=0.0, prefix_matching=False):
         # type: (Sequence[str], Sequence[str], bool, float) -> List[InfoType]
 
         """
@@ -112,6 +122,8 @@ def GetInfos(self, shard_ids=(), keys=(), allow_partial=False, max_age=0.0):
             will raise an exception.
         :param max_age: If specified and non-zero, only shard infos whose age is less-
             than-or-equal-to max_age will be returned in the response.
+        :param prefix_matching: If True, we treat each key in `keys` as a prefix of a
+            field in `full_info`
         """
 
         resp = []
@@ -131,7 +143,7 @@ def GetInfos(self, shard_ids=(), keys=(), allow_partial=False, max_age=0.0):
                 if max_age and info_age > max_age:
                     logger.debug('Shard %s info age %s > %s (max-age); Skipping', shard_id, info_age, max_age)
                     continue
-                msg = self._filter_info(full_info=shard.info, keys=set(keys))
+                msg = self._filter_info(full_info=shard.info, keys=set(keys), prefix_matching=prefix_matching)
                 msg['meta']['info_age'] = info_age
             except KeyError as e:
                 if allow_partial:

tests/test_infoServicer.py

@@ -117,8 +117,29 @@ def test_query_filtering(self):
 
         for k in resp_info.keys():
             self.assertNotIn('removed', k, 'key in response INFO that should have been filtered')
+        # `resp_info` includes the returned INFO fields + an internal-use field called `meta`
         self.assertEqual(len(resp_info), 3, 'Unexpected keys in response ({})'.format(resp_info.keys()))
 
+    def test_query_filtering_not_exact(self):
+        shard_id = 'shard-1'
+        info = {
+            'dummy_key1': 'dummy',
+            'dummy_key2': 'dummy',
+            'asdf_key1': 'dummy',
+            'asdf_key2': 'dummy',
+            'removed_key1': 'removed',
+            'removed_key2': 'removed',
+        }
+
+        ShardPublisher.add_shard(self.make_shard(shard_id, info=info))
+
+        resp_info = self.servicer.GetInfos(shard_ids=[shard_id], keys=['dummy', 'asdf'], prefix_matching=True)[0]
+
+        for k in resp_info.keys():
+            self.assertNotIn('removed', k, 'key in response INFO that should have been filtered')
+        # `resp_info` includes the returned INFO fields + an internal-use field called `meta`
+        self.assertEqual(len(resp_info), 5, 'Unexpected keys in response ({})'.format(resp_info.keys()))
+
     def test_query_missing(self):
         shard_ids = ['shard-1', 'shard-2']