Search enhancements

Author: mahewitt

README.rst

@@ -10,7 +10,14 @@ Note: This is currently a WIP so will be subject to breaking changes.
 
 Change Log
 ==========
-                  
+                   
+0.0.9
+-----
+
+- Search query_by_kind function
+- Search query_by_id supports limit
+- Search query supports a specific query
+                   
 0.0.8
 -----
 

src/osdu/__init__.py

@@ -4,4 +4,4 @@
 # license information.
 # -----------------------------------------------------------------------------
 
-__VERSION__ = "0.0.8"
+__VERSION__ = "0.0.9"

src/osdu/search/_client.py

@@ -61,16 +61,24 @@ def query_all_aggregated(self) -> dict:
         response_json = self._client.post_returning_json(self.api_url("query"), request_data)
         return response_json
 
-    def query(self, kind: str = None, identifier: str = None, limit: int = None) -> dict:
+    def query(
+        self, kind: str = None, identifier: str = None, query: str = None, limit: int = None
+    ) -> dict:
         """Query records
 
         Args:
             kind (str): kind to query for
             identifier (str): id to query for
+            query (str): a specific query
+            limit (str): limit on number of records to return
 
         Returns:
             dict: containing the result
         """
+
+        if identifier is not None and query is not None:
+            raise ValueError("You can't specify both identifier and query")
+
         request_data = {}
         if kind is None:
             request_data["kind"] = "*:*:*:*"
@@ -80,13 +88,16 @@ def query(self, kind: str = None, identifier: str = None, limit: int = None) ->
         if identifier is not None:
             request_data["query"] = f'id:("{identifier}")'
 
+        if query is not None:
+            request_data["query"] = query
+
         if limit is not None:
             request_data["limit"] = limit
 
         response_json = self._client.post_returning_json(self.api_url("query"), request_data)
         return response_json
 
-    def query_by_id(self, identifier: str) -> dict:
+    def query_by_id(self, identifier: str, limit: int = None) -> dict:
         """Returns a list of all kinds including number of records
 
         Args:
@@ -100,5 +111,25 @@ def query_by_id(self, identifier: str) -> dict:
             "query": f'id:("{identifier}")',
         }
 
+        if limit is not None:
+            request_data["limit"] = limit
+
+        response_json = self._client.post_returning_json(self.api_url("query"), request_data)
+        return response_json
+
+    def query_by_kind(self, kind: str, limit: int = None) -> dict:
+        """Returns a list of all records for the given kind
+
+        Args:
+            kind (str): kind to query for
+
+        Returns:
+            dict: containing the result
+        """
+        request_data = {"kind": kind}
+
+        if limit is not None:
+            request_data["limit"] = limit
+
         response_json = self._client.post_returning_json(self.api_url("query"), request_data)
         return response_json

tests/search/test_search.py

@@ -114,14 +114,25 @@ def test_query_all_aggregated_http_error(self, _):
 
     # region test query
     @params(
-        ("osdu:wks:dataset--File.Generic:1.0.0", None, None),
-        ("*:*:*:*", None, None),
-        (None, None, None),
-        (None, "opendes:reference-data--ResourceSecurityClassification:RESTRICTED", None),
-        (None, "opendes:reference-data--ResourceSecurityClassification:RESTRICTED", 20),
-        ("*:*:*:*", "opendes:reference-data--ResourceSecurityClassification:RESTRICTED", None),
+        ("osdu:wks:dataset--File.Generic:1.0.0", None, None, None),
+        ("*:*:*:*", None, None, None),
+        (None, None, None, None),
+        (None, "opendes:reference-data--ResourceSecurityClassification:RESTRICTED", None, None),
+        (None, "opendes:reference-data--ResourceSecurityClassification:RESTRICTED", None, 20),
+        (
+            None,
+            None,
+            'data.WellboreID:("opendes:master-data--Wellbore:ad215042-05db-2b7e-e053-c818a488c79a")',
+            None,
+        ),
+        (
+            "*:*:*:*",
+            "opendes:reference-data--ResourceSecurityClassification:RESTRICTED",
+            None,
+            None,
+        ),
     )
-    def test_query(self, kind, identifier, limit):
+    def test_query(self, kind, identifier, query, limit):
         """Test the query function"""
         request_data = {}
 
@@ -133,6 +144,9 @@ def test_query(self, kind, identifier, limit):
         if identifier is not None:
             request_data["query"] = f'id:("{identifier}")'
 
+        if query is not None:
+            request_data["query"] = query
+
         if limit is not None:
             request_data["limit"] = limit
 
@@ -150,14 +164,22 @@ def test_query(self, kind, identifier, limit):
             client = create_dummy_client()
             search_client = SearchClient(client)
 
-            response_data = search_client.query(kind, identifier, limit)
+            response_data = search_client.query(kind, identifier, query, limit)
 
             mock_post_returning_json.assert_called_once()
             mock_post_returning_json.assert_called_with(
                 "http://www.test.com/api/search/v2/query", request_data
             )
             self.assertEqual(expected_response_data, response_data)
 
+    def test_query_cant_pass_id_and_query(self):
+        """Test query with id and query faile"""
+        with self.assertRaises(ValueError):
+            client = create_dummy_client()
+            search_client = SearchClient(client)
+
+            _ = search_client.query(identifier="id", query="query")
+
     @mock.patch.object(OsduClient, "post_returning_json", side_effect=HTTPError(1))
     def test_query_http_error(self, _):
         """Test query http errors are propogated"""
@@ -170,13 +192,47 @@ def test_query_http_error(self, _):
     # endregion test query
 
     # region test query_by_id
+    @params(
+        ("id1", None),
+        ("id1", 2),
+        ("id", None),
+        ("id", 2),
+    )
+    def test_query_by_id(self, identifier, limit):
+        """Test the query_by_id function"""
+        request_data = {
+            "kind": "*:*:*:*",
+            "query": f'id:("{identifier}")',
+        }
+        if limit is not None:
+            request_data["limit"] = limit
+
+        expected_response_data = {
+            "results": [],
+            "totalCount": 1,
+        }
+        with mock.patch(
+            "osdu.client.OsduClient.post_returning_json", return_value=expected_response_data
+        ) as mock_post_returning_json:
+            client = create_dummy_client()
+            search_client = SearchClient(client)
+
+            response_data = search_client.query_by_id(identifier, limit)
+
+            mock_post_returning_json.assert_called_once()
+            mock_post_returning_json.assert_called_with(
+                "http://www.test.com/api/search/v2/query", request_data
+            )
+            self.assertEqual(expected_response_data, response_data)
+
     @params("id1", "id")
-    def test_query_by_id(self, identifier):
+    def test_query_by_id_defaults(self, identifier):
         """Test the query_by_id function"""
         request_data = {
             "kind": "*:*:*:*",
             "query": f'id:("{identifier}")',
         }
+
         expected_response_data = {
             "results": [],
             "totalCount": 1,
@@ -206,6 +262,75 @@ def test_query_by_id_http_error(self, _):
 
     # endregion test query_by_id
 
+    # region test query_by_kind
+    @params(
+        ("kind1", None),
+        ("kind1", 2),
+        ("kind", None),
+        ("kind", 2),
+    )
+    def test_query_by_kind(self, kind, limit):
+        """Test the query_by_kind function"""
+        request_data = {
+            "kind": kind,
+        }
+        if limit is not None:
+            request_data["limit"] = limit
+
+        expected_response_data = {
+            "results": [],
+            "totalCount": 1,
+        }
+        with mock.patch(
+            "osdu.client.OsduClient.post_returning_json", return_value=expected_response_data
+        ) as mock_post_returning_json:
+            client = create_dummy_client()
+            search_client = SearchClient(client)
+
+            response_data = search_client.query_by_kind(kind, limit)
+
+            mock_post_returning_json.assert_called_once()
+            mock_post_returning_json.assert_called_with(
+                "http://www.test.com/api/search/v2/query", request_data
+            )
+            self.assertEqual(expected_response_data, response_data)
+
+    @params("kind1", "kind")
+    def test_query_by_kind_defaults(self, kind):
+        """Test the query_by_kind function"""
+        request_data = {
+            "kind": kind,
+        }
+
+        expected_response_data = {
+            "results": [],
+            "totalCount": 1,
+        }
+        with mock.patch(
+            "osdu.client.OsduClient.post_returning_json", return_value=expected_response_data
+        ) as mock_post_returning_json:
+            client = create_dummy_client()
+            search_client = SearchClient(client)
+
+            response_data = search_client.query_by_kind(kind)
+
+            mock_post_returning_json.assert_called_once()
+            mock_post_returning_json.assert_called_with(
+                "http://www.test.com/api/search/v2/query", request_data
+            )
+            self.assertEqual(expected_response_data, response_data)
+
+    @mock.patch.object(OsduClient, "post_returning_json", side_effect=HTTPError(1))
+    def test_query_by_kind_http_error(self, _):
+        """Test query_by_kind http errors are propogated"""
+        with self.assertRaises(HTTPError):
+            client = create_dummy_client()
+            search_client = SearchClient(client)
+
+            _ = search_client.query_by_kind("kind")
+
+    # endregion test query_by_kind
+
 
 if __name__ == "__main__":
     import nose2