From 9d6ab4c40f231d7f805228820756d28da936c91d Mon Sep 17 00:00:00 2001 From: Surya Avala Date: Mon, 2 Nov 2020 03:28:51 +1100 Subject: [PATCH] fix(proxy): fixes GCP inverse proxy url priority. Fixes #4284 (#4702) * fix: fixes proxy url priority uses proxy url that is closest to the given zone Fixes #4284 * style: formatting proxy with yapf using .style.yapf provided in the repo --- proxy/get_proxy_url.py | 124 +++++++++++++++++++----------------- proxy/get_proxy_url_test.py | 32 +++++++--- 2 files changed, 88 insertions(+), 68 deletions(-) diff --git a/proxy/get_proxy_url.py b/proxy/get_proxy_url.py index a802242885..760cd7a8ee 100644 --- a/proxy/get_proxy_url.py +++ b/proxy/get_proxy_url.py @@ -12,7 +12,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """CLI tool that returns URL of the proxy for particular zone and version.""" import argparse import functools @@ -22,12 +21,13 @@ import re import requests try: - unicode + unicode except NameError: - unicode = str + unicode = str + def urls_for_zone(zone, location_to_urls_map): - """Returns list of potential proxy URLs for a given zone. + """Returns list of potential proxy URLs for a given zone. Returns: List of possible URLs, in order of proximity. @@ -41,71 +41,79 @@ def urls_for_zone(zone, location_to_urls_map): ... } """ - zone_match = re.match("((([a-z]+)-[a-z]+)\d+)-[a-z]", zone) - if not zone_match: - raise ValueError("Incorrect zone specified: {}".format(zone)) + zone_match = re.match("((([a-z]+)-[a-z]+)\d+)-[a-z]", zone) + if not zone_match: + raise ValueError("Incorrect zone specified: {}".format(zone)) - # e.g. zone = us-west1-b - region = zone_match.group(1) # us-west1 - approx_region = zone_match.group(2) # us-west - country = zone_match.group(3) # us + # e.g. zone = us-west1-b + region = zone_match.group(1) # us-west1 + approx_region = zone_match.group(2) # us-west + country = zone_match.group(3) # us - urls = [] - if region in location_to_urls_map: - urls.extend(location_to_urls_map[region]) + urls = [] + if region in location_to_urls_map: + urls.extend([ + url for url in location_to_urls_map[region] if url not in urls + ]) - region_regex = re.compile("([a-z]+-[a-z]+)\d+") - for location in location_to_urls_map: - region_match = region_regex.match(location) - if region_match and region_match.group(1) == approx_region: - urls.extend(location_to_urls_map[location]) + region_regex = re.compile("([a-z]+-[a-z]+)\d+") + for location in location_to_urls_map: + region_match = region_regex.match(location) + if region_match and region_match.group(1) == approx_region: + urls.extend([ + url for url in location_to_urls_map[location] if url not in urls + ]) - if country in location_to_urls_map: - urls.extend(location_to_urls_map[country]) + if country in location_to_urls_map: + urls.extend([ + url for url in location_to_urls_map[country] if url not in urls + ]) - return set(urls) + return urls def main(): - unicode_type = functools.partial(unicode, encoding="utf8") - parser = argparse.ArgumentParser( - description="Get proxy URL") - parser.add_argument("--config-file-path", required=True, type=unicode_type) - parser.add_argument("--location", required=True, type=unicode_type) - parser.add_argument("--version", required=True, type=unicode_type) + unicode_type = functools.partial(unicode, encoding="utf8") + parser = argparse.ArgumentParser(description="Get proxy URL") + parser.add_argument("--config-file-path", required=True, type=unicode_type) + parser.add_argument("--location", required=True, type=unicode_type) + parser.add_argument("--version", required=True, type=unicode_type) - args = parser.parse_args() - with open(args.config_file_path, "r") as config_file: - data = json.loads(config_file.read()) + args = parser.parse_args() + with open(args.config_file_path, "r") as config_file: + data = json.loads(config_file.read()) - agent_containers_config = data["agent-docker-containers"] - version = args.version - if version not in agent_containers_config: - version = "latest" - if version not in agent_containers_config: - raise ValueError("Version latest not found in the config file.") - container_config = agent_containers_config[version] - regional_urls = container_config["proxy-urls"] + agent_containers_config = data["agent-docker-containers"] + version = args.version + if version not in agent_containers_config: + version = "latest" + if version not in agent_containers_config: + raise ValueError("Version latest not found in the config file.") + container_config = agent_containers_config[version] + regional_urls = container_config["proxy-urls"] - location = args.location - urls = urls_for_zone(location, regional_urls) - if not urls: - raise ValueError("No valid URLs found for zone: {}".format(location)) + location = args.location + urls = urls_for_zone(location, regional_urls) + if not urls: + raise ValueError("No valid URLs found for zone: {}".format(location)) + + for url in urls: + try: + status_code = requests.head(url).status_code + except requests.ConnectionError: + pass + expected_codes = frozenset([307]) + # 307 - Temporary Redirect, Proxy server sends this if VM has access rights. + if status_code in expected_codes: + logging.debug("Status code from the url %s", status_code) + print(url) + exit(0) + logging.debug( + "Incorrect status_code from the server: %s. Expected: %s", + status_code, expected_codes + ) + raise ValueError("No working URL found") - for url in urls: - try: - status_code = requests.head(url).status_code - except requests.ConnectionError: - pass - expected_codes = frozenset([307]) - # 307 - Temporary Redirect, Proxy server sends this if VM has access rights. - if status_code in expected_codes: - logging.debug("Status code from the url %s", status_code) - print(url) - exit(0) - logging.debug("Incorrect status_code from the server: %s. Expected: %s", - status_code, expected_codes) - raise ValueError("No working URL found") if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/proxy/get_proxy_url_test.py b/proxy/get_proxy_url_test.py index 7f3f73b99c..e23e5403a1 100644 --- a/proxy/get_proxy_url_test.py +++ b/proxy/get_proxy_url_test.py @@ -22,24 +22,36 @@ url_map_json = """ { "us": ["https://datalab-us-west1.cloud.google.com"], "us-west1": ["https://datalab-us-west1.cloud.google.com"], + "us-west2": ["https://datalab-us-west2.cloud.google.com"], "us-east1": ["https://datalab-us-east1.cloud.google.com"] } """ + class TestUrlsForZone(unittest.TestCase): - def test_get_urls(self): - self.assertEqual( - set(["https://datalab-us-east1.cloud.google.com","https://datalab-us-west1.cloud.google.com"]), - urls_for_zone("us-east1-a",json.loads(url_map_json))) + def test_get_urls(self): + self.assertEqual([ + "https://datalab-us-east1.cloud.google.com", + "https://datalab-us-west1.cloud.google.com" + ], urls_for_zone("us-east1-a", json.loads(url_map_json))) + def test_get_urls_no_match(self): + self.assertEqual([], + urls_for_zone( + "euro-west1-a", json.loads(url_map_json) + )) - def test_get_urls_no_match(self): - self.assertEqual(set([]), urls_for_zone("euro-west1-a",json.loads(url_map_json))) + def test_get_urls_incorrect_format(self): + with self.assertRaises(ValueError): + urls_for_zone("weird-format-a", json.loads(url_map_json)) + + def test_get_urls_priority(self): + self.assertEqual([ + "https://datalab-us-west1.cloud.google.com", + "https://datalab-us-west2.cloud.google.com" + ], urls_for_zone("us-west1-a", json.loads(url_map_json))) - def test_get_urls_incorrect_format(self): - with self.assertRaises(ValueError): - urls_for_zone("weird-format-a",json.loads(url_map_json)) if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main()