From 9fd59f78806a5c3ef7da576b06bc40055938ae3a Mon Sep 17 00:00:00 2001
From: Abhishek Raut <rauta@vmware.com>
Date: Thu, 12 Jan 2017 04:45:07 -0800
Subject: [PATCH] Add support to search resources based on tags or resource
 type

This patch adds a new util method to the NsxLib class.
NsxLib will expose a search method to retrieve objects from
backend based on their tags and resource type. Tags argument
must be present in order to search.
Tags are supplied in the following form:
    [{'scope': <scope_val>, 'tag': <tag_value>}, ...]

Change-Id: I304e9c44e55657e652b2a8236e85602c295cf22b
---
 vmware_nsxlib/tests/unit/v3/test_resources.py | 35 +++++++++++++++++++
 vmware_nsxlib/v3/__init__.py                  | 33 +++++++++++++++++
 vmware_nsxlib/v3/exceptions.py                |  4 +++
 3 files changed, 72 insertions(+)

diff --git a/vmware_nsxlib/tests/unit/v3/test_resources.py b/vmware_nsxlib/tests/unit/v3/test_resources.py
index aa89982a..5a22ed75 100644
--- a/vmware_nsxlib/tests/unit/v3/test_resources.py
+++ b/vmware_nsxlib/tests/unit/v3/test_resources.py
@@ -777,3 +777,38 @@ class IpPoolTestCase(nsxlib_testcase.NsxClientTestCase):
         test_client.assert_json_call(
             'get', pool,
             'https://1.2.3.4/api/v1/pools/ip-pools/%s/allocations' % uuid)
+
+
+class TestNsxSearch(nsxlib_testcase.NsxClientTestCase):
+
+    def test_nsx_search_tags(self):
+        """Test search of resources with the specified tag."""
+        with mock.patch.object(self.nsxlib.client, 'url_get') as search:
+            user_tags = [{'scope': 'user', 'tag': 'k8s'}]
+            query = self.nsxlib._build_query(tags=user_tags)
+            self.nsxlib.search_by_tags(tags=user_tags)
+            search.assert_called_with('search?query=%s' % query)
+
+    def test_nsx_search_tags_and_resource_type(self):
+        """Test search of specified resource with the specified tag."""
+        with mock.patch.object(self.nsxlib.client, 'url_get') as search:
+            user_tags = [{'scope': 'user', 'tag': 'k8s'}]
+            res_type = 'LogicalPort'
+            query = self.nsxlib._build_query(tags=user_tags)
+            # Add resource_type to the query
+            query = "%s AND %s" % (res_type, query)
+            self.nsxlib.search_by_tags(tags=user_tags, resource_type=res_type)
+            search.assert_called_with('search?query=%s' % query)
+
+    def test_nsx_search_invalid_query_fail(self):
+        """Test search query failure for missing tag argument."""
+        self.assertRaises(exceptions.NsxSearchInvalidQuery,
+                          self.nsxlib.search_by_tags,
+                          tags=None, resource_type=None)
+
+    def test_nsx_search_invalid_tags_fail(self):
+        """Test search of resources with the invalid tag."""
+        user_tags = [{'scope': 'user', 'invalid_tag_key': 'k8s'}]
+        self.assertRaises(exceptions.NsxSearchInvalidQuery,
+                          self.nsxlib._build_query,
+                          tags=user_tags)
diff --git a/vmware_nsxlib/v3/__init__.py b/vmware_nsxlib/v3/__init__.py
index 9d583853..9cf3ea04 100644
--- a/vmware_nsxlib/v3/__init__.py
+++ b/vmware_nsxlib/v3/__init__.py
@@ -99,6 +99,39 @@ class NsxLib(object):
     def subscribe(self, callback, event):
         self.cluster.subscribe(callback, event)
 
+    # TODO(abhiraut): Revisit this method to generate complex boolean
+    #                 queries to search resources.
+    def search_by_tags(self, tags, resource_type=None):
+        """Return the list of resources searched based on tags.
+
+        Currently the query only supports AND boolean operator.
+        :param tags: List of dictionaries containing tags. Each
+                     NSX tag dictionary is of the form:
+                     {'scope': <scope_key>, 'tag': <tag_value>}
+        :param resource_type: Optional string parameter to limit the
+                              scope of the search to the given ResourceType.
+        """
+        if not tags:
+            reason = _("Missing required argument 'tags'")
+            raise exceptions.NsxSearchInvalidQuery(reason=reason)
+        # Query will return nothing if the same scope is repeated.
+        query_tags = self._build_query(tags)
+        query = resource_type
+        if query:
+            query += " AND %s" % query_tags
+        else:
+            query = query_tags
+        url = "search?query=%s" % query
+        return self.client.url_get(url)
+
+    def _build_query(self, tags):
+        try:
+            return " AND ".join(['tags.scope:%(scope)s AND '
+                                 'tags.tag:%(tag)s' % item for item in tags])
+        except KeyError as e:
+            reason = _('Missing key:%s in tags') % str(e)
+            raise exceptions.NsxSearchInvalidQuery(reason=reason)
+
 
 class NsxLibPortMirror(utils.NsxLibApiBase):
 
diff --git a/vmware_nsxlib/v3/exceptions.py b/vmware_nsxlib/v3/exceptions.py
index 0e674810..5fb0bf23 100644
--- a/vmware_nsxlib/v3/exceptions.py
+++ b/vmware_nsxlib/v3/exceptions.py
@@ -119,3 +119,7 @@ class NumberOfNsgroupCriteriaTagsReached(ManagerError):
 class SecurityGroupMaximumCapacityReached(ManagerError):
     message = _("Security Group %(sg_id)s has reached its maximum capacity, "
                 "no more ports can be associated with this security-group.")
+
+
+class NsxSearchInvalidQuery(NsxLibException):
+    message = _("Invalid input for NSX search query. Reason: %(reason)s")