From 24f29505eb8566a2030a456b8791dac8e7897be2 Mon Sep 17 00:00:00 2001
From: BinamB <numb.baj@gmail.com>
Date: Mon, 13 May 2024 10:36:56 -0500
Subject: [PATCH] Add parameter to init endpoint

---
 fence/blueprints/data/blueprint.py | 11 ++++---
 fence/blueprints/data/indexd.py    | 22 +++++++++++++-
 tests/data/test_blank_index.py     | 46 ++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/fence/blueprints/data/blueprint.py b/fence/blueprints/data/blueprint.py
index 5dc1055c8..ab07b2886 100755
--- a/fence/blueprints/data/blueprint.py
+++ b/fence/blueprints/data/blueprint.py
@@ -199,16 +199,19 @@ def upload_data_file():
 def init_multipart_upload():
     """
     Initialize a multipart upload request
-
-    NOTE This endpoint does not currently accept a `bucket` parameter like
-    `POST /upload` and `GET /upload/<GUID>` do.
     """
     params = flask.request.get_json()
     if not params:
         raise UserError("wrong Content-Type; expected application/json")
     if "file_name" not in params:
         raise UserError("missing required argument `file_name`")
-    blank_index = BlankIndex(file_name=params["file_name"])
+
+    create_record = params.get("create_record", True)
+
+    guid = params.get("did")
+    blank_index = BlankIndex(
+        file_name=params["file_name"], guid=guid, create_record=create_record
+    )
 
     default_expires_in = flask.current_app.config.get("MAX_PRESIGNED_URL_TTL", 3600)
     expires_in = get_valid_expiration(
diff --git a/fence/blueprints/data/indexd.py b/fence/blueprints/data/indexd.py
index f7b9488f6..8a6b36856 100755
--- a/fence/blueprints/data/indexd.py
+++ b/fence/blueprints/data/indexd.py
@@ -208,6 +208,7 @@ def _get_client_id():
 
     return client_id
 
+
 def prepare_presigned_url_audit_log(protocol, indexed_file):
     """
     Store in `flask.g.audit_data` the data needed to record an audit log.
@@ -232,7 +233,13 @@ class BlankIndex(object):
     """
 
     def __init__(
-        self, uploader=None, file_name=None, logger_=None, guid=None, authz=None
+        self,
+        uploader=None,
+        file_name=None,
+        logger_=None,
+        guid=None,
+        authz=None,
+        create_record=True,
     ):
         self.logger = logger_ or logger
         self.indexd = (
@@ -253,7 +260,10 @@ def __init__(
         self.file_name = file_name
         self.authz = authz
 
+        self.create_record = create_record
+
         # if a guid is not provided, this will create a blank record for you
+        self.guid = None
         self.guid = guid or self.index_document["did"]
 
     @cached_property
@@ -266,6 +276,16 @@ def index_document(self):
                 response from indexd (the contents of the record), containing ``guid``
                 and ``url``
         """
+
+        # if the record already exists in indexd, just fetch the record
+        if not self.create_record:
+            index_url = self.indexd.rstrip("/") + "/index/index/" + self.guid
+            indexd_response = requests.get(index_url)
+            if indexd_response.status_code == 200:
+                self.logger.info("found record with guid: {}".format(self.guid))
+                data = indexd_response.json()
+                return data
+
         index_url = self.indexd.rstrip("/") + "/index/blank/"
         params = {"uploader": self.uploader, "file_name": self.file_name}
 
diff --git a/tests/data/test_blank_index.py b/tests/data/test_blank_index.py
index 7586d85be..a43ff82d0 100755
--- a/tests/data/test_blank_index.py
+++ b/tests/data/test_blank_index.py
@@ -507,3 +507,49 @@ def test_make_signed_url_missing_configuration_key(app, indexd_client):
                 signed_url = blank_index.make_signed_url(
                     file_name="some file name", protocol=indexed_file_location
                 )
+
+
+def test_blank_index_upload_to_specific_guid(
+    app, client, auth_client, encoded_creds_jwt, user_client
+):
+    """
+    test BlankIndex upload
+    POST /data/upload
+    """
+    did = str(uuid.uuid4())
+    data_requests_mocker = mock.patch(
+        "fence.blueprints.data.indexd.requests", new_callable=mock.Mock
+    )
+    arborist_requests_mocker = mock.patch(
+        "gen3authz.client.arborist.client.httpx.Client.request", new_callable=mock.Mock
+    )
+    with data_requests_mocker as data_requests, arborist_requests_mocker as arborist_requests:
+        data_requests.post.return_value = MockResponse(
+            {
+                "did": str(uuid.uuid4()),
+                "rev": str(uuid.uuid4())[:8],
+                "baseid": str(uuid.uuid4()),
+            }
+        )
+        data_requests.post.return_value.status_code = 200
+        arborist_requests.return_value = MockResponse({"auth": True})
+        arborist_requests.return_value.status_code = 200
+        headers = {
+            "Authorization": "Bearer " + encoded_creds_jwt.jwt,
+            "Content-Type": "application/json",
+        }
+        file_name = "asdf"
+        data = json.dumps({"file_name": file_name})
+        response = client.post("/data/upload", headers=headers, data=data)
+        indexd_url = app.config.get("INDEXD") or app.config.get("BASE_URL") + "/index"
+        endpoint = indexd_url + "/index/blank/"
+        indexd_auth = (config["INDEXD_USERNAME"], config["INDEXD_PASSWORD"])
+        data_requests.post.assert_called_once_with(
+            endpoint,
+            auth=indexd_auth,
+            json={"file_name": file_name, "uploader": user_client.username},
+            headers={},
+        )
+        assert response.status_code == 201, response.json
+        assert "guid" in response.json
+        assert "url" in response.json