support list for select argument (Azure#14991)

xiangyan99 · web-flow · commit be079df27e16 · 2020-11-05T10:31:47.000-08:00
* support list for select argument

* update
diff --git a/sdk/search/azure-search-documents/azure/search/documents/_internal/_queries.py b/sdk/search/azure-search-documents/azure/search/documents/_internal/_queries.py
@@ -99,3 +99,16 @@ class SuggestQuery(_QueryBase):
     _request_type = SuggestRequest
 
     __doc__ = SuggestRequest.__doc__
+
+    def select(self, *fields):
+        # type: (*str) -> None
+        """Update the `select` property for the search results.
+
+        :param fields: An list of fields for the query result to return.
+        :type expression: str
+        :raises: ValueError
+        """
+        if not fields:
+            raise ValueError("At least one field must be provided")
+
+        self._request.select = ",".join(fields)
diff --git a/sdk/search/azure-search-documents/azure/search/documents/_internal/_search_client.py b/sdk/search/azure-search-documents/azure/search/documents/_internal/_search_client.py
@@ -4,7 +4,6 @@
 # license information.
 # --------------------------------------------------------------------------
 from typing import cast, List, TYPE_CHECKING
-
 import six
 
 from azure.core.tracing.decorator import distributed_trace
@@ -261,10 +260,12 @@ def search(self, search_text, **kwargs):
             scoring_profile=scoring_profile,
             search_fields=search_fields,
             search_mode=search_mode,
-            select=select,
+            select=select if isinstance(select, six.string_types) else None,
             skip=skip,
             top=top
         )
+        if isinstance(select, list):
+            query.select(select)
 
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         return SearchItemPaged(
@@ -336,10 +337,11 @@ def suggest(self, search_text, suggester_name, **kwargs):
             minimum_coverage=minimum_coverage,
             order_by=order_by,
             search_fields=search_fields,
-            select=select,
+            select=select if isinstance(select, six.string_types) else None,
             top=top
         )
-
+        if isinstance(select, list):
+            query.select(select)
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         response = self._client.documents.suggest_post(
             suggest_request=query.request, **kwargs
diff --git a/sdk/search/azure-search-documents/azure/search/documents/_internal/aio/_search_client_async.py b/sdk/search/azure-search-documents/azure/search/documents/_internal/aio/_search_client_async.py
@@ -4,6 +4,7 @@
 # license information.
 # --------------------------------------------------------------------------
 from typing import cast, List, TYPE_CHECKING
+import six
 
 from azure.core.tracing.decorator_async import distributed_trace_async
 from ._paging import AsyncSearchItemPaged, AsyncSearchPageIterator
@@ -223,11 +224,12 @@ async def search(self, search_text, **kwargs):
             scoring_profile=scoring_profile,
             search_fields=search_fields,
             search_mode=search_mode,
-            select=select,
+            select=select if isinstance(select, six.string_types) else None,
             skip=skip,
             top=top
         )
-
+        if isinstance(select, list):
+            query.select(select)
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         return AsyncSearchItemPaged(
             self._client, query, kwargs, page_iterator_class=AsyncSearchPageIterator
@@ -298,10 +300,11 @@ async def suggest(self, search_text, suggester_name, **kwargs):
             minimum_coverage=minimum_coverage,
             order_by=order_by,
             search_fields=search_fields,
-            select=select,
+            select=select if isinstance(select, six.string_types) else None,
             top=top
         )
-
+        if isinstance(select, list):
+            query.select(select)
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         response = await self._client.documents.suggest_post(
             suggest_request=query.request, **kwargs
diff --git a/sdk/search/azure-search-documents/tests/async_tests/recordings/test_search_client_search_live_async.test_get_search_filter_array.yaml b/sdk/search/azure-search-documents/tests/async_tests/recordings/test_search_client_search_live_async.test_get_search_filter_array.yaml
@@ -0,0 +1,43 @@
+interactions:
+- request:
+    body: '{"filter": "category eq ''Budget''", "orderby": "hotelName desc", "search":
+      "WiFi", "select": "hotelName,category,description"}'
+    headers:
+      Accept:
+      - application/json;odata.metadata=none
+      Content-Length:
+      - '125'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-search-documents/11.1.0b2 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+    method: POST
+    uri: https://search6e521a9a.search.windows.net/indexes('drgqefsg')/docs/search.post.search?api-version=2020-06-30
+  response:
+    body:
+      string: '{"value":[{"@search.score":1.0225849,"hotelName":"Express Rooms","description":"Pretty
+        good hotel","category":"Budget"},{"@search.score":1.0225849,"hotelName":"EconoStay","description":"Very
+        popular hotel in town","category":"Budget"},{"@search.score":2.3832736,"hotelName":"Countryside
+        Hotel","description":"Save up to 50% off traditional hotels.  Free WiFi, great
+        location near downtown, full kitchen, washer & dryer, 24/7 support, bowling
+        alley, fitness center and more.","category":"Budget"},{"@search.score":1.0225849,"hotelName":"Comfy
+        Place","description":"Another good hotel","category":"Budget"}]}'
+    headers:
+      cache-control: no-cache
+      content-encoding: gzip
+      content-length: '442'
+      content-type: application/json; odata.metadata=none
+      date: Fri, 28 Aug 2020 21:07:31 GMT
+      elapsed-time: '73'
+      expires: '-1'
+      odata-version: '4.0'
+      pragma: no-cache
+      preference-applied: odata.include-annotations="*"
+      request-id: 7c9fee41-e972-11ea-9d02-5cf37071153c
+      strict-transport-security: max-age=15724800; includeSubDomains
+      vary: Accept-Encoding
+    status:
+      code: 200
+      message: OK
+    url: https://search6e521a9a.search.windows.net/indexes('drgqefsg')/docs/search.post.search?api-version=2020-06-30
+version: 1
diff --git a/sdk/search/azure-search-documents/tests/async_tests/recordings/test_search_client_search_live_async.test_get_search_simple_paged.yaml b/sdk/search/azure-search-documents/tests/async_tests/recordings/test_search_client_search_live_async.test_get_search_simple_paged.yaml
@@ -0,0 +1,132 @@
+interactions:
+- request:
+    body: '{"search": "hotel"}'
+    headers:
+      Accept:
+      - application/json;odata.metadata=none
+      Content-Length:
+      - '19'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-search-documents/11.1.0b2 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+    method: POST
+    uri: https://search6e8a1a80.search.windows.net/indexes('drgqefsg')/docs/search.post.search?api-version=2020-06-30
+  response:
+    body:
+      string: '{"value":[{"@search.score":2.245087,"hotelId":"10","hotelName":"Countryside
+        Hotel","description":"Save up to 50% off traditional hotels.  Free WiFi, great
+        location near downtown, full kitchen, washer & dryer, 24/7 support, bowling
+        alley, fitness center and more.","descriptionFr":"\u00c9conomisez jusqu''\u00e0
+        50% sur les h\u00f4tels traditionnels.  WiFi gratuit, tr\u00e8s bien situ\u00e9
+        pr\u00e8s du centre-ville, cuisine compl\u00e8te, laveuse & s\u00e9cheuse,
+        support 24/7, bowling, centre de fitness et plus encore.","category":"Budget","tags":["24-hour
+        front desk service","coffee in lobby","restaurant"],"parkingIncluded":false,"smokingAllowed":true,"lastRenovationDate":"1999-09-06T00:00:00Z","rating":3,"location":{"type":"Point","coordinates":[-78.940483,35.90416],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":{"streetAddress":"6910
+        Fayetteville Rd","city":"Durham","stateProvince":"NC","country":"USA","postalCode":"27713"},"rooms":[{"description":"Suite,
+        1 King Bed (Amenities)","descriptionFr":"Suite, 1 tr\u00e8s grand lit (Services)","type":"Suite","baseRate":2.44,"bedOptions":"1
+        King Bed","sleepsCount":2,"smokingAllowed":true,"tags":["coffee maker"]},{"description":"Budget
+        Room, 1 Queen Bed (Amenities)","descriptionFr":"Chambre \u00c9conomique, 1
+        grand lit (Services)","type":"Budget Room","baseRate":7.69,"bedOptions":"1
+        Queen Bed","sleepsCount":2,"smokingAllowed":false,"tags":["coffee maker"]}]},{"@search.score":1.778081,"hotelId":"3","hotelName":"EconoStay","description":"Very
+        popular hotel in town","descriptionFr":"H\u00f4tel le plus populaire en ville","category":"Budget","tags":["wifi","budget"],"parkingIncluded":true,"smokingAllowed":false,"lastRenovationDate":"1995-07-01T00:00:00Z","rating":4,"location":{"type":"Point","coordinates":[-122.131577,46.678581],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":null,"rooms":[]},{"@search.score":1.778081,"hotelId":"4","hotelName":"Express
+        Rooms","description":"Pretty good hotel","descriptionFr":"Assez bon h\u00f4tel","category":"Budget","tags":["wifi","budget"],"parkingIncluded":true,"smokingAllowed":false,"lastRenovationDate":"1995-07-01T00:00:00Z","rating":4,"location":{"type":"Point","coordinates":[-122.131577,48.678581],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":null,"rooms":[]},{"@search.score":1.778081,"hotelId":"5","hotelName":"Comfy
+        Place","description":"Another good hotel","descriptionFr":"Un autre bon h\u00f4tel","category":"Budget","tags":["wifi","budget"],"parkingIncluded":true,"smokingAllowed":false,"lastRenovationDate":"2012-08-12T00:00:00Z","rating":4,"location":{"type":"Point","coordinates":[-122.131577,48.678581],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":null,"rooms":[]},{"@search.score":1.5749159,"hotelId":"2","hotelName":"Roach
+        Motel","description":"Cheapest hotel in town. Infact, a motel.","descriptionFr":"H\u00f4tel
+        le moins cher en ville. Infact, un motel.","category":"Budget","tags":["motel","budget"],"parkingIncluded":true,"smokingAllowed":true,"lastRenovationDate":"1982-04-28T00:00:00Z","rating":1,"location":{"type":"Point","coordinates":[-122.131577,49.678581],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":null,"rooms":[]},{"@search.score":1.4634744,"hotelId":"1","hotelName":"Fancy
+        Stay","description":"Best hotel in town if you like luxury hotels. They have
+        an amazing infinity pool, a spa, and a really helpful concierge. The location
+        is perfect -- right downtown, close to all the tourist attractions. We highly
+        recommend this hotel.","descriptionFr":"Meilleur h\u00f4tel en ville si vous
+        aimez les h\u00f4tels de luxe. Ils ont une magnifique piscine \u00e0 d\u00e9bordement,
+        un spa et un concierge tr\u00e8s utile. L''emplacement est parfait \u2013
+        en plein centre, \u00e0 proximit\u00e9 de toutes les attractions touristiques.
+        Nous recommandons fortement cet h\u00f4tel.","category":"Luxury","tags":["pool","view","wifi","concierge"],"parkingIncluded":false,"smokingAllowed":false,"lastRenovationDate":"2010-06-27T00:00:00Z","rating":5,"location":{"type":"Point","coordinates":[-122.131577,47.678581],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":null,"rooms":[]},{"@search.score":0.5496142,"hotelId":"9","hotelName":"Secret
+        Point Motel","description":"The hotel is ideally located on the main commercial
+        artery of the city in the heart of New York. A few minutes away is Time''s
+        Square and the historic centre of the city, as well as other places of interest
+        that make New York one of America''s most attractive and cosmopolitan cities.","descriptionFr":"L''h\u00f4tel
+        est id\u00e9alement situ\u00e9 sur la principale art\u00e8re commerciale de
+        la ville en plein c\u0153ur de New York. A quelques minutes se trouve la place
+        du temps et le centre historique de la ville, ainsi que d''autres lieux d''int\u00e9r\u00eat
+        qui font de New York l''une des villes les plus attractives et cosmopolites
+        de l''Am\u00e9rique.","category":"Boutique","tags":["pool","air conditioning","concierge"],"parkingIncluded":false,"smokingAllowed":true,"lastRenovationDate":"1970-01-18T05:00:00Z","rating":4,"location":{"type":"Point","coordinates":[-73.975403,40.760586],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":{"streetAddress":"677
+        5th Ave","city":"New York","stateProvince":"NY","country":"USA","postalCode":"10022"},"rooms":[{"description":"Budget
+        Room, 1 Queen Bed (Cityside)","descriptionFr":"Chambre \u00c9conomique, 1
+        grand lit (c\u00f4t\u00e9 ville)","type":"Budget Room","baseRate":9.69,"bedOptions":"1
+        Queen Bed","sleepsCount":2,"smokingAllowed":true,"tags":["vcr/dvd"]},{"description":"Budget
+        Room, 1 King Bed (Mountain View)","descriptionFr":"Chambre \u00c9conomique,
+        1 tr\u00e8s grand lit (Mountain View)","type":"Budget Room","baseRate":8.09,"bedOptions":"1
+        King Bed","sleepsCount":2,"smokingAllowed":true,"tags":["vcr/dvd","jacuzzi
+        tub"]}]}]}'
+    headers:
+      cache-control: no-cache
+      content-encoding: gzip
+      content-length: '2378'
+      content-type: application/json; odata.metadata=none
+      date: Fri, 28 Aug 2020 21:07:53 GMT
+      elapsed-time: '83'
+      expires: '-1'
+      odata-version: '4.0'
+      pragma: no-cache
+      preference-applied: odata.include-annotations="*"
+      request-id: 8a56c06f-e972-11ea-a554-5cf37071153c
+      strict-transport-security: max-age=15724800; includeSubDomains
+      vary: Accept-Encoding
+    status:
+      code: 200
+      message: OK
+    url: https://search6e8a1a80.search.windows.net/indexes('drgqefsg')/docs/search.post.search?api-version=2020-06-30
+- request:
+    body: '{"search": "motel"}'
+    headers:
+      Accept:
+      - application/json;odata.metadata=none
+      Content-Length:
+      - '19'
+      Content-Type:
+      - application/json
+      User-Agent:
+      - azsdk-python-search-documents/11.1.0b2 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+    method: POST
+    uri: https://search6e8a1a80.search.windows.net/indexes('drgqefsg')/docs/search.post.search?api-version=2020-06-30
+  response:
+    body:
+      string: '{"value":[{"@search.score":8.376183,"hotelId":"2","hotelName":"Roach
+        Motel","description":"Cheapest hotel in town. Infact, a motel.","descriptionFr":"H\u00f4tel
+        le moins cher en ville. Infact, un motel.","category":"Budget","tags":["motel","budget"],"parkingIncluded":true,"smokingAllowed":true,"lastRenovationDate":"1982-04-28T00:00:00Z","rating":1,"location":{"type":"Point","coordinates":[-122.131577,49.678581],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":null,"rooms":[]},{"@search.score":0.8858137,"hotelId":"9","hotelName":"Secret
+        Point Motel","description":"The hotel is ideally located on the main commercial
+        artery of the city in the heart of New York. A few minutes away is Time''s
+        Square and the historic centre of the city, as well as other places of interest
+        that make New York one of America''s most attractive and cosmopolitan cities.","descriptionFr":"L''h\u00f4tel
+        est id\u00e9alement situ\u00e9 sur la principale art\u00e8re commerciale de
+        la ville en plein c\u0153ur de New York. A quelques minutes se trouve la place
+        du temps et le centre historique de la ville, ainsi que d''autres lieux d''int\u00e9r\u00eat
+        qui font de New York l''une des villes les plus attractives et cosmopolites
+        de l''Am\u00e9rique.","category":"Boutique","tags":["pool","air conditioning","concierge"],"parkingIncluded":false,"smokingAllowed":true,"lastRenovationDate":"1970-01-18T05:00:00Z","rating":4,"location":{"type":"Point","coordinates":[-73.975403,40.760586],"crs":{"type":"name","properties":{"name":"EPSG:4326"}}},"address":{"streetAddress":"677
+        5th Ave","city":"New York","stateProvince":"NY","country":"USA","postalCode":"10022"},"rooms":[{"description":"Budget
+        Room, 1 Queen Bed (Cityside)","descriptionFr":"Chambre \u00c9conomique, 1
+        grand lit (c\u00f4t\u00e9 ville)","type":"Budget Room","baseRate":9.69,"bedOptions":"1
+        Queen Bed","sleepsCount":2,"smokingAllowed":true,"tags":["vcr/dvd"]},{"description":"Budget
+        Room, 1 King Bed (Mountain View)","descriptionFr":"Chambre \u00c9conomique,
+        1 tr\u00e8s grand lit (Mountain View)","type":"Budget Room","baseRate":8.09,"bedOptions":"1
+        King Bed","sleepsCount":2,"smokingAllowed":true,"tags":["vcr/dvd","jacuzzi
+        tub"]}]}]}'
+    headers:
+      cache-control: no-cache
+      content-encoding: gzip
+      content-length: '1269'
+      content-type: application/json; odata.metadata=none
+      date: Fri, 28 Aug 2020 21:07:55 GMT
+      elapsed-time: '7'
+      expires: '-1'
+      odata-version: '4.0'
+      pragma: no-cache
+      preference-applied: odata.include-annotations="*"
+      request-id: 8ac67395-e972-11ea-9c35-5cf37071153c
+      strict-transport-security: max-age=15724800; includeSubDomains
+      vary: Accept-Encoding
+    status:
+      code: 200
+      message: OK
+    url: https://search6e8a1a80.search.windows.net/indexes('drgqefsg')/docs/search.post.search?api-version=2020-06-30
+version: 1
diff --git a/sdk/search/azure-search-documents/tests/async_tests/test_search_client_search_live_async.py b/sdk/search/azure-search-documents/tests/async_tests/test_search_client_search_live_async.py
@@ -58,6 +58,23 @@ async def test_get_search_simple(self, api_key, endpoint, index_name, **kwargs):
                 results.append(x)
             assert len(results) == 2
 
+    @ResourceGroupPreparer(random_name_enabled=True)
+    @SearchServicePreparer(schema=SCHEMA, index_batch=BATCH)
+    async def test_get_search_simple_paged(self, api_key, endpoint, index_name, **kwargs):
+        client = SearchClient(
+            endpoint, index_name, AzureKeyCredential(api_key)
+        )
+        async with client:
+            results = []
+            async for x in await client.search(search_text="hotel", top=3):
+                results.append(x)
+            assert len(results) == 7
+
+            results = []
+            async for x in await client.search(search_text="motel"):
+                results.append(x)
+            assert len(results) == 2
+
     @ResourceGroupPreparer(random_name_enabled=True)
     @SearchServicePreparer(schema=SCHEMA, index_batch=BATCH)
     async def test_get_search_filter(self, api_key, endpoint, index_name, **kwargs):
@@ -88,6 +105,36 @@ async def test_get_search_filter(self, api_key, endpoint, index_name, **kwargs):
             assert all(set(x) == expected for x in results)
             assert all(x["category"] == "Budget" for x in results)
 
+    @ResourceGroupPreparer(random_name_enabled=True)
+    @SearchServicePreparer(schema=SCHEMA, index_batch=BATCH)
+    async def test_get_search_filter_array(self, api_key, endpoint, index_name, **kwargs):
+        client = SearchClient(
+            endpoint, index_name, AzureKeyCredential(api_key)
+        )
+
+        async with client:
+            results = []
+            select = ("hotelName", "category", "description")
+            async for x in await client.search(
+                    search_text="WiFi",
+                    filter="category eq 'Budget'",
+                    select=select,
+                    order_by="hotelName desc"
+            ):
+                results.append(x)
+            assert [x["hotelName"] for x in results] == sorted(
+                [x["hotelName"] for x in results], reverse=True
+            )
+            expected = {
+                "category",
+                "hotelName",
+                "description",
+                "@search.score",
+                "@search.highlights",
+            }
+            assert all(set(x) == expected for x in results)
+            assert all(x["category"] == "Budget" for x in results)
+
     @ResourceGroupPreparer(random_name_enabled=True)
     @SearchServicePreparer(schema=SCHEMA, index_batch=BATCH)
     async def test_get_search_counts(self, api_key, endpoint, index_name, **kwargs):
diff --git a/sdk/search/azure-search-documents/tests/recordings/test_search_client_search_live.test_get_search_filter_array.yaml b/sdk/search/azure-search-documents/tests/recordings/test_search_client_search_live.test_get_search_filter_array.yaml
diff --git a/sdk/search/azure-search-documents/tests/recordings/test_search_client_search_live.test_get_search_simple_paged.yaml b/sdk/search/azure-search-documents/tests/recordings/test_search_client_search_live.test_get_search_simple_paged.yaml
diff --git a/sdk/search/azure-search-documents/tests/test_search_client_search_live.py b/sdk/search/azure-search-documents/tests/test_search_client_search_live.py
