PR-5069 Refactor locate endpoint

requirements/requirements.in

@@ -1,6 +1,5 @@
 cachetools
 cfnresponse
 chalice
-flatdict
 git+https://github.com/asfadmin/rain-api-core.git@318aac226c92cf6f60cc8821d6d94669485972c6
 netaddr

requirements/requirements.txt

@@ -22,8 +22,6 @@ click==8.1.7
     # via chalice
 cryptography==41.0.7
     # via pyjwt
-flatdict==4.0.1
-    # via -r requirements/requirements.in
 inquirer==2.10.1
     # via chalice
 jinja2==3.1.2

tests/data/old_style_bucket_map_example.yaml

@@ -0,0 +1,33 @@
+RAW:
+  PLATFORM-A:         pa-raw
+  PLATFORM-B:         pb-raw
+DATA-TYPE-1:
+  PLATFORM-A:         pa-dt1
+  PLATFORM-B:         pb-dt1
+BROWSE:
+  PLATFORM-A:         pa-bro
+  PLATFORM-B:         pb-bro
+PRIVATE:
+  PLATFORM-A:         pa-priv
+  PLATFORM-B:         pb-priv
+HEADERS:
+  BROWSE:
+    bucket: pa-bro
+    headers:
+      custom-header-1: custom-header-1-value
+      custom-header-2: custom-header-2-value
+  PRIVATE:
+    bucket: pa-priv
+    headers:
+      custom-header-3: custom-header-3-value
+      custom-header-4: custom-header-4-value
+
+PUBLIC_BUCKETS:
+  - pa-bro
+  - pb-bro
+
+PRIVATE_BUCKETS:
+  pa-priv:
+    - PRIVATE_GROUP_A
+  pb-priv:
+    - PRIVATE_GROUP_B

tests/test_app.py

@@ -569,6 +569,7 @@ def test_get_user_ip(current_request):
 @mock.patch(f"{MODULE}.get_role_creds", autospec=True)
 @mock.patch(f"{MODULE}.get_role_session", autospec=True)
 @mock.patch(f"{MODULE}.get_presigned_url", autospec=True)
+@mock.patch(f"{MODULE}.b_map", None)
 def test_try_download_from_bucket(
     mock_get_presigned_url,
     mock_get_role_session,
@@ -587,7 +588,6 @@ def test_try_download_from_bucket(
     client = mock_get_role_session().client()
     client.get_bucket_location.return_value = {"LocationConstraint": "us-east-1"}
     client.head_object.return_value = {"ContentLength": 2048}
-    app.b_map = None
 
     response = app.try_download_from_bucket("somebucket", "somefile", user_profile, {})
     client.head_object.assert_called_once()
@@ -901,22 +901,86 @@ def test_version(mock_retrieve_secret, monkeypatch, client):
 
 
 @mock.patch(f"{MODULE}.get_yaml_file", autospec=True)
-def test_locate(mock_get_yaml_file, mock_retrieve_secret, data_path, client):
+@mock.patch(f"{MODULE}.b_map", None)
+def test_locate(
+    mock_get_yaml_file,
+    mock_retrieve_secret,
+    data_path,
+    monkeypatch,
+    client,
+):
     del mock_retrieve_secret
 
     with open(data_path / "bucket_map_example.yaml") as f:
         mock_get_yaml_file.return_value = yaml.full_load(f)
 
+    monkeypatch.setenv("BUCKETNAME_PREFIX", "")
+
     response = client.http.get("/locate?bucket_name=pa-dt1")
+    assert response.status_code == 200
     assert response.json_body == ["DATA-TYPE-1/PLATFORM-A"]
+
+    response = client.http.get("/locate?bucket_name=nonexistent")
+    assert response.status_code == 404
+    assert response.body == b"No route defined for nonexistent"
+
+
+@mock.patch(f"{MODULE}.get_yaml_file", autospec=True)
+@mock.patch(f"{MODULE}.b_map", None)
+def test_locate_old_style_bucket_map(
+    mock_get_yaml_file,
+    mock_retrieve_secret,
+    data_path,
+    monkeypatch,
+    client,
+):
+    del mock_retrieve_secret
+
+    with open(data_path / "old_style_bucket_map_example.yaml") as f:
+        mock_get_yaml_file.return_value = yaml.full_load(f)
+
+    monkeypatch.setenv("BUCKETNAME_PREFIX", "")
+
+    response = client.http.get("/locate?bucket_name=pa-dt1")
     assert response.status_code == 200
+    assert response.json_body == ["DATA-TYPE-1/PLATFORM-A"]
 
     response = client.http.get("/locate?bucket_name=nonexistent")
+    assert response.status_code == 404
     assert response.body == b"No route defined for nonexistent"
+
+
+@mock.patch(f"{MODULE}.get_yaml_file", autospec=True)
+@mock.patch(f"{MODULE}.b_map", None)
+def test_locate_bucket_name_prefix(
+    mock_get_yaml_file,
+    mock_retrieve_secret,
+    data_path,
+    monkeypatch,
+    client,
+):
+    del mock_retrieve_secret
+
+    with open(data_path / "bucket_map_example.yaml") as f:
+        mock_get_yaml_file.return_value = yaml.full_load(f)
+
+    monkeypatch.setenv("BUCKETNAME_PREFIX", "bucket-prefix-")
+
+    response = client.http.get("/locate?bucket_name=bucket-prefix-pa-dt1")
+    assert response.status_code == 200
+    assert response.json_body == ["DATA-TYPE-1/PLATFORM-A"]
+
+    response = client.http.get("/locate?bucket_name=pa-dt1")
     assert response.status_code == 404
+    assert response.body == b"No route defined for pa-dt1"
+
+    response = client.http.get("/locate?bucket_name=nonexistent")
+    assert response.status_code == 404
+    assert response.body == b"No route defined for nonexistent"
 
 
 @pytest.mark.parametrize("req", ("/locate", "/locate?foo=bar"))
+@mock.patch(f"{MODULE}.b_map", None)
 def test_locate_missing_bucket(mock_retrieve_secret, client, req):
     del mock_retrieve_secret
 
@@ -929,41 +993,6 @@ def test_locate_missing_bucket(mock_retrieve_secret, client, req):
     }
 
 
-def test_collapse_bucket_configuration():
-    bucket_map = {
-        "foo": "bar",
-        "key1": {
-            "key2": {
-                "bucket": "bucket1"
-            }
-        },
-        "bucket": {
-            "bucket": "bucket2"
-        },
-        "key3": {
-            "bucket": {
-                "bucket": {
-                    "bucket": "bucket3"
-                }
-            }
-        }
-    }
-    app.collapse_bucket_configuration(bucket_map)
-
-    assert bucket_map == {
-        "foo": "bar",
-        "key1": {
-            "key2": "bucket1"
-        },
-        "bucket": "bucket2",
-        "key3": {
-            "bucket": {
-                "bucket": "bucket3"
-            }
-        }
-    }
-
-
 def test_get_range_header_val(current_request):
     current_request.headers = {"Range": "v1"}
     assert app.get_range_header_val() == "v1"
@@ -1152,7 +1181,6 @@ def test_dynamic_url(
 
     mock_get_profile.return_value = user_profile
     current_request.uri_params = {"proxy": "DATA-TYPE-1/PLATFORM-A/OBJECT_1"}
-    app.b_map = None
 
     # Can't use the chalice test client here as it doesn't seem to understand the `{proxy+}` route
     response = app.dynamic_url()
@@ -1445,6 +1473,7 @@ def test_dynamic_url_directory(
 @mock.patch(f"{MODULE}.RequestAuthorizer._handle_auth_bearer_header", autospec=True)
 @mock.patch(f"{MODULE}.JwtManager.get_header_to_set_auth_cookie", autospec=True)
 @mock.patch(f"{MODULE}.JWT_COOKIE_NAME", "asf-cookie")
+@mock.patch(f"{MODULE}.b_map", None)
 def test_dynamic_url_bearer_auth(
     mock_get_header_to_set_auth_cookie,
     mock_handle_auth_bearer_header,

thin_egress_app/app.py

@@ -13,7 +13,6 @@
 import boto3
 import cachetools
 import chalice
-import flatdict
 from botocore.config import Config as bc_Config
 from botocore.exceptions import ClientError
 from cachetools.func import ttl_cache
@@ -720,16 +719,30 @@ def version():
 @app.route('/locate')
 @with_trace(context={})
 def locate():
+    timer = Timer()
+    timer.mark('restore_bucket_vars()')
+
     query_params = app.current_request.query_params
     if query_params is None or query_params.get('bucket_name') is None:
         return Response(body='Required "bucket_name" query paramater not specified',
                         status_code=400,
                         headers={'Content-Type': 'text/plain'})
+
+    restore_bucket_vars()
+    timer.mark()
+
     bucket_name = query_params.get('bucket_name')
-    bucket_map = collapse_bucket_configuration(get_yaml_file(conf_bucket, bucket_map_file)['MAP'])
-    search_map = flatdict.FlatDict(bucket_map, delimiter='/')
-    matching_paths = [key for key, value in search_map.items() if value == bucket_name]
-    if (len(matching_paths) > 0):
+    matching_paths = [
+        entry.bucket_path
+        for entry in b_map.entries()
+        if entry.bucket == bucket_name
+    ]
+    log.debug('matching_paths: %s', matching_paths)
+
+    log.debug('timing for locate(): ')
+    timer.log_all(log)
+
+    if matching_paths:
         return Response(body=json.dumps(matching_paths),
                         status_code=200,
                         headers={'Content-Type': 'application/json'})
@@ -738,17 +751,6 @@ def locate():
                     headers={'Content-Type': 'text/plain'})
 
 
-@with_trace()
-def collapse_bucket_configuration(bucket_map):
-    for k, v in bucket_map.items():
-        if isinstance(v, dict):
-            if 'bucket' in v:
-                bucket_map[k] = v['bucket']
-            else:
-                collapse_bucket_configuration(v)
-    return bucket_map
-
-
 @with_trace()
 def get_range_header_val():
     if 'Range' in app.current_request.headers: