diff --git a/rbac/api/common/pagination.py b/rbac/api/common/pagination.py index 4747d28ea..82d41eb52 100644 --- a/rbac/api/common/pagination.py +++ b/rbac/api/common/pagination.py @@ -17,13 +17,13 @@ """Common pagination class.""" import logging +import re +from urllib.parse import urlparse from rest_framework.pagination import LimitOffsetPagination from rest_framework.response import Response from rest_framework.utils.urls import replace_query_param -from api import API_VERSION - PATH_INFO = "PATH_INFO" logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -38,16 +38,13 @@ class StandardResultsSetPagination(LimitOffsetPagination): def link_rewrite(request, link): """Rewrite the link based on the path header to only provide partial url.""" url = link - version = "v{}/".format(API_VERSION) if PATH_INFO in request.META: - try: - local_api_index = link.index(version) - path = request.META.get(PATH_INFO) - path_api_index = path.index(version) - path_link = "{}{}" - url = path_link.format(path[:path_api_index], link[local_api_index:]) - except ValueError: - logger.warning('Unable to rewrite link as "{}" was not found.'.format(version)) + url_components = urlparse(link) + path_and_query = url_components.path + (f"?{url_components.query}" if url_components.query else "") + if bool(re.search("/v[0-9]/", path_and_query)): + url = path_and_query + else: + logger.warning(f"Unable to rewrite link as no version was not found in {path_and_query}.") return url def get_first_link(self): diff --git a/tests/api/common/test_pagination.py b/tests/api/common/test_pagination.py index d5ad8e652..54b743484 100644 --- a/tests/api/common/test_pagination.py +++ b/tests/api/common/test_pagination.py @@ -34,6 +34,24 @@ def test_link_rewrite(self): result = StandardResultsSetPagination.link_rewrite(request, link) self.assertEqual(expected, result) + def test_link_rewrite_rbac_prefix(self): + """Test the link rewrite for RBAC prefix.""" + request = Mock() + request.META = {PATH_INFO: "/api/rbac/v1/roles/"} + link = "http://localhost:8000/api/rbac/v1/roles/?limit=10&offset=0" + expected = "/api/rbac/v1/roles/?limit=10&offset=0" + result = StandardResultsSetPagination.link_rewrite(request, link) + self.assertEqual(expected, result) + + def test_link_rewrite_rbac_prefix_v2(self): + """Test the link rewrite for RBAC prefix for v2 APIs.""" + request = Mock() + request.META = {PATH_INFO: "/api/rbac/v2/workspaces/"} + link = "http://localhost:8000/api/rbac/v2/workspaces/?limit=10&offset=0" + expected = "/api/rbac/v2/workspaces/?limit=10&offset=0" + result = StandardResultsSetPagination.link_rewrite(request, link) + self.assertEqual(expected, result) + def test_link_rewrite_err(self): """Test the link rewrite.""" request = Mock()