Add function for "safe" NSX search

This change adds the search_all_by_tags_safe function, which
operaties exactly as search_all_by_tags but raises an exception
after the first page of data is returned, it the search resultset
size exceeds a specific threshold.

This enables better control over cases where a search query selects
too many records and cannot be fully served by NSX search.

Change-Id: I475377bd6f6a8e78a5a92c369d9d0e1c95ea4829

vmware_nsxlib/tests/unit/v3/test_resources.py

@@ -2008,16 +2008,61 @@ class TestNsxSearch(nsxlib_testcase.NsxClientTestCase):
         self.assertEqual(3, len(results))
         self.assertEqual([{"id": "s1"}, {"id": "s2"}, {"id": "s3"}], results)
 
+    def test_nsx_search_all_safe_max_rs_ok(self):
+        """Test search all base method."""
+        mock_search = mock.Mock()
+        mock_search.side_effect = [
+            {"cursor": "2",
+             "result_count": 3,
+             "results": [{"id": "s1"}, {"id": "s2"}]},
+            {"cursor": "3",
+             "result_count": 3,
+             "results": [{"id": "s3"}]}]
+        args = "foo"
+        kwargs = {"a1": "v1", "a2": "v2"}
+        results = self.nsxlib._search_all(mock_search,
+                                          safe_mode=True,
+                                          max_rs=10000,
+                                          *args, **kwargs)
+        mock_search.assert_has_calls([
+            mock.call(*args, cursor=0, page_size=None, **kwargs),
+            mock.call(*args, cursor=2, page_size=None, **kwargs)])
+        self.assertEqual(3, len(results))
+        self.assertEqual([{"id": "s1"}, {"id": "s2"}, {"id": "s3"}], results)
+
+    def test_nsx_search_all_safe_max_rs_too_large(self):
+        """Test search all base method."""
+        mock_search = mock.Mock()
+        # We don't really need to put 40000 items in the
+        # result, this way of mocking the response is ok
+        mock_search.side_effect = [
+            {"cursor": "2",
+             "result_count": "40000",
+             "results": [{"id": "s1"}, {"id": "s2"}]},
+            {"cursor": "3",
+             "result_count": "40000",
+             "results": [{"id": "s3"}]}]
+        args = "foo"
+        kwargs = {"a1": "v1", "a2": "v2"}
+        with self.assertRaises(exceptions.SearchResultsetTooLarge) as exc:
+            self.nsxlib._search_all(
+                mock_search, safe_mode=True, max_rs=10000,
+                *args, **kwargs)
+            self.assertEqual(40000, exc.result_count)
+        # We expect to raise after the 1st call to search API
+        mock_search.assert_has_calls(
+            [mock.call(*args, cursor=0, page_size=None, **kwargs)])
+
     def test_nsx_search_all_by_tags(self):
         """Test search all of resources with the specified tag."""
         with mock.patch.object(self.nsxlib.client, 'url_get') as search:
             search.side_effect = [
                 {"cursor": "2",
-                 "result_count": 3,
+                 "result_count": "3",
                  "results": [{"id": "s1"},
                              {"id": "s2"}]},
                 {"cursor": "3",
-                 "result_count": 3,
+                 "result_count": "3",
                  "results": [{"id": "s3"}]}]
             user_tags = [{'scope': 'user', 'tag': 'k8s'}]
             query = self.nsxlib._build_query(tags=user_tags)
@@ -2028,6 +2073,45 @@ class TestNsxSearch(nsxlib_testcase.NsxClientTestCase):
                           silent=False)])
             self.assertEqual(3, len(results))
 
+    def test_nsx_search_all_by_tags_safe_max_rs_ok(self):
+        with mock.patch.object(self.nsxlib.client, 'url_get') as search:
+            search.side_effect = [
+                {"cursor": "2",
+                 "result_count": "3",
+                 "results": [{"id": "s1"},
+                             {"id": "s2"}]},
+                {"cursor": "3",
+                 "result_count": "3",
+                 "results": [{"id": "s3"}]}]
+            user_tags = [{'scope': 'user', 'tag': 'k8s'}]
+            query = self.nsxlib._build_query(tags=user_tags)
+            results = self.nsxlib.search_all_by_tags_safe(
+                10000, tags=user_tags)
+            search.assert_has_calls([
+                mock.call(self.search_path % query, silent=False),
+                mock.call((self.search_path + '&cursor=2') % query,
+                          silent=False)])
+            self.assertEqual(3, len(results))
+
+    def test_nsx_search_all_by_tags_safe_max_rs_too_large(self):
+        with mock.patch.object(self.nsxlib.client, 'url_get') as search:
+            search.side_effect = [
+                {"cursor": "2",
+                 "result_count": "40000",
+                 "results": [{"id": "s1"},
+                             {"id": "s2"}]},
+                {"cursor": "3",
+                 "result_count": "40000",
+                 "results": [{"id": "s3"}]}]
+            user_tags = [{'scope': 'user', 'tag': 'k8s'}]
+            query = self.nsxlib._build_query(tags=user_tags)
+            with self.assertRaises(exceptions.SearchResultsetTooLarge) as exc:
+                self.nsxlib.search_all_by_tags_safe(
+                    10000, tags=user_tags)
+                self.asseerEqual(40000, exc.result_count)
+            search.assert_has_calls([
+                mock.call(self.search_path % query, silent=False)])
+
     @mock.patch("vmware_nsxlib.v3.lib.NsxLibBase._search_all")
     def test_nsx_search_all_by_attribute_values(self, mock_search_all):
         """Test search all resources with the specified attribute values."""

vmware_nsxlib/v3/exceptions.py

@@ -66,6 +66,14 @@ class NsxLibInvalidInput(NsxLibException):
     message = _("Invalid input for operation: %(error_message)s.")
 
 
+class SearchResultsetTooLarge(NsxLibException):
+    message = _("Search query returns too many results: %(result_count)d")
+
+    def __init__(self, **kwargs):
+        super(SearchResultsetTooLarge, self).__init__(**kwargs)
+        self.result_count = kwargs.get('result_count')
+
+
 class ManagerError(NsxLibException):
     message = _("Unexpected error from backend manager (%(manager)s) "
                 "for %(operation)s%(details)s")

vmware_nsxlib/v3/lib.py

@@ -270,6 +270,9 @@ class NsxLibBase(object, metaclass=abc.ABCMeta):
         results = []
         cursor = 0
         page_size = None
+        # The following kwargs must not be passed to search_func
+        safe_mode = kwargs.pop('safe_mode', False)
+        max_rs = kwargs.pop('max_rs', nsx_constants.SEARCH_MAX_RS)
         while True:
             try:
                 response = search_func(*args, cursor=cursor,
@@ -287,6 +290,11 @@ class NsxLibBase(object, metaclass=abc.ABCMeta):
             results.extend(response['results'])
             cursor = int(response['cursor'])
             result_count = int(response['result_count'])
+            # If we want a safe search, raise an exception if the
+            # search is returning too many results
+            if safe_mode and result_count > max_rs:
+                raise exceptions.SearchResultsetTooLarge(
+                    result_count=result_count)
             if cursor >= result_count:
                 return results
 
@@ -296,6 +304,15 @@ class NsxLibBase(object, metaclass=abc.ABCMeta):
                                 resource_type=resource_type, tags=tags,
                                 **extra_attrs)
 
+    def search_all_by_tags_safe(self, max_rs, tags,
+                                resource_type=None, **extra_attrs):
+        return self._search_all(self.search_by_tags,
+                                safe_mode=True,
+                                max_rs=max_rs,
+                                resource_type=resource_type,
+                                tags=tags,
+                                **extra_attrs)
+
     def search_all_resource_by_attributes(self, resource_type, **attributes):
         """Return all resources of a given type matching specific attributes.
 

vmware_nsxlib/v3/nsx_constants.py

@@ -120,6 +120,8 @@ NUM_ALLOWED_IP_ADDRESSES = 128
 NUM_ALLOWED_IP_ADDRESSES_v4 = NUM_ALLOWED_IP_ADDRESSES
 NUM_ALLOWED_IP_ADDRESSES_v6 = 15
 MAX_STATIC_ROUTES = 26
+# Max number of entries returned by a search query
+SEARCH_MAX_RS = 60000
 
 # QoS directions egress/ingress
 EGRESS = 'egress'