[ID-323] Allow tnu to use either Martha or DRSHub to resolve DRS URIs (#398)

* make drs resolver configurable

* add env var to readme

* fix lint error

Author: Shakespeared

Changes.md

@@ -1,3 +1,7 @@
+# Changes for v0.11.0 (2022-11-09)
+Added support for calling DRSHub for DRS resolution
+Made DRS Resolver configurable by setting `DRS_RESOLVER` to `martha` or `drshub`
+
 # Changes for v0.10.0 (2022-06-30)
 
 # Changes for v0.9.0 (2022-06-30)
@@ -128,4 +132,4 @@ Config override for tests (#72)
 ## 0.2.1 - 06/12/2020
 - Added requester pays set-up to drs module
 - Added ENV variable for local testing to __init__
-- Added some print messages to drs module to let user know what is happening
+- Added some print messages to drs module to let user know what is happening

README.md

@@ -240,7 +240,7 @@ This will run tests against Terra and Martha Dev using Jade Dev DRS url (make su
 This will run tests against Terra and Martha Prod (make sure you have proper access to DRS urls, workspace and Google bucket)
 
 3. log in with your Google credentials using `gcloud auth application-default login` with your Terra Prod account
-4. set `export GOOGLE_PROJECT=firecloud-cgl; export TERRA_DEPLOYMENT_ENV=prod`
+4. set `export GOOGLE_PROJECT=firecloud-cgl; export TERRA_DEPLOYMENT_ENV=prod; export TNU_BLOBSTORE_TEST_GS_BUCKET=tnu-test-bucket`
 5. run in package root:
     - `make test`: skips controlled and dev access tests
     - `make controlled_access_test`: runs tests marked as `controlled_access`

terra_notebook_utils/__init__.py

@@ -5,6 +5,7 @@
 WORKSPACE_GOOGLE_PROJECT = os.environ.get('GOOGLE_PROJECT')  # This env var is set in Terra Cloud Environments
 TERRA_DEPLOYMENT_ENV = os.environ.get('TERRA_DEPLOYMENT_ENV', 'prod')
 MARTHA_URL_VERSION = os.environ.get('MARTHA_URL_VERSION', 'martha_v3')
+DRS_RESOLVER = os.environ.get('DRS_RESOLVER', 'martha')  # values for DRS_RESOLVER are either 'drshub' or 'martha'
 
 if not WORKSPACE_GOOGLE_PROJECT:
     WORKSPACE_GOOGLE_PROJECT = os.environ.get('GCP_PROJECT')  # Useful for running outside of notebook
@@ -20,3 +21,8 @@
 IO_CONCURRENCY = 3
 
 MARTHA_URL = f"https://us-central1-broad-dsde-{TERRA_DEPLOYMENT_ENV}.cloudfunctions.net/{MARTHA_URL_VERSION}"
+DRSHUB_URL = f"https://drshub.dsde-{TERRA_DEPLOYMENT_ENV}.broadinstitute.org/api/v4/drs/resolve"
+if DRS_RESOLVER == "martha":
+    DRS_RESOLVER_URL = MARTHA_URL
+else:
+    DRS_RESOLVER_URL = DRSHUB_URL

terra_notebook_utils/drs.py

@@ -7,7 +7,7 @@
 from typing import Dict, List, Tuple, Optional, Union, Iterable
 from requests import Response
 
-from terra_notebook_utils import WORKSPACE_BUCKET, WORKSPACE_NAME, MARTHA_URL, WORKSPACE_NAMESPACE, \
+from terra_notebook_utils import WORKSPACE_BUCKET, WORKSPACE_NAME, DRS_RESOLVER_URL, WORKSPACE_NAMESPACE, \
     WORKSPACE_GOOGLE_PROJECT
 from terra_notebook_utils import workspace, gs, tar_gz, TERRA_DEPLOYMENT_ENV, _GS_SCHEMA
 from terra_notebook_utils.utils import is_notebook
@@ -60,18 +60,18 @@ def enable_requester_pays(workspace_name: Optional[str]=WORKSPACE_NAME,
                        "You will not be able to access DRS URIs that interact with requester pays buckets.")
 
 def get_drs(drs_url: str, fields: List[str]) -> Response:
-    """Request DRS information from martha."""
+    """Request DRS information from DRS Resolver."""
     access_token = gs.get_access_token()
 
     headers = {
         'authorization': f"Bearer {access_token}",
         'content-type': "application/json"
     }
 
-    logger.debug(f"Resolving DRS uri '{drs_url}' through '{MARTHA_URL}'.")
+    logger.debug(f"Resolving DRS uri '{drs_url}' through '{DRS_RESOLVER_URL}'.")
 
     json_body = dict(url=drs_url, fields=fields)
-    resp = http.post(MARTHA_URL, headers=headers, json=json_body)
+    resp = http.post(DRS_RESOLVER_URL, headers=headers, json=json_body)
 
     if 200 != resp.status_code:
         logger.warning(resp.content)
@@ -99,18 +99,18 @@ def access(drs_url: str,
     info = get_drs_info(drs_url, access_url=True)
 
     if info.access_url:
-        martha_url = info.access_url.strip()
-        response = requests.get(martha_url, headers={'Range': 'bytes=0-1'})
+        drs_resolver_url = info.access_url.strip()
+        response = requests.get(drs_resolver_url, headers={'Range': 'bytes=0-1'})
         if response.status_code < 300:
-            return martha_url
+            return drs_resolver_url
         logger.warning('Got an invalid/inaccessible signed URL... attempting to generate an alternate signed URL...')
 
     url = gs.get_signed_url(bucket=info.bucket_name,
                             key=info.key,
                             sa_credentials=info.credentials)
     # attempt to get the first byte; we'll get an HTTPError error if we need requester pays
-    # TODO: hopefully Martha returns this information eventually and we can avoid this check,
-    #  but even in the meantime, there's probably a better way of doing this
+    # TODO: hopefully the DRS Resolver will return this information eventually and
+    #  we can avoid this check, but even in the meantime, there's probably a better way of doing this
     response = requests.get(url, headers={'Range': 'bytes=0-1'})
     if b'Bucket is a requester pays bucket' in response.content and response.status_code >= 400:
         url = gs.get_signed_url(bucket=info.bucket_name,
@@ -164,8 +164,8 @@ def _drs_info_from_martha_v2(drs_url: str, drs_data: dict) -> DRSInfo:
     else:
         raise DRSResolutionError(f"No metadata was returned for DRS uri '{drs_url}'")
 
-def _drs_info_from_martha_v3(drs_url: str, drs_data: dict) -> DRSInfo:
-    """Convert response from martha_v3 to DRSInfo."""
+def _drs_info_from_martha_v3_or_drshub(drs_url: str, drs_data: dict) -> DRSInfo:
+    """Convert DRS response to DRSInfo."""
     access_url: Optional[str] = None
     if drs_data.get('accessUrl') is not None:
         access_url = drs_data['accessUrl'].get('url')
@@ -193,14 +193,14 @@ def get_drs_info(drs_url: str, access_url: bool = False) -> DRSInfo:
               "name",
               "timeUpdated",
               "googleServiceAccount"]
-    # only include this field if necessary because it significantly increases the time for Martha to respond
+    # only include this field if necessary because it significantly increases the time for the DRS Resolver to respond
     if access_url:
         fields.append("accessUrl")
     drs_data = get_drs(drs_url, fields).json()
     if 'dos' in drs_data:
         return _drs_info_from_martha_v2(drs_url, drs_data)
     else:
-        return _drs_info_from_martha_v3(drs_url, drs_data)
+        return _drs_info_from_martha_v3_or_drshub(drs_url, drs_data)
 
 def get_drs_blob(drs_url_or_info: Union[str, DRSInfo],
                  billing_project: Optional[str]=WORKSPACE_GOOGLE_PROJECT) -> Union[GSBlob, URLBlob]:

terra_notebook_utils/gs.py

@@ -82,10 +82,10 @@ def get_signed_url(bucket: str,
     """
     default_service_account_credentials = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
 
-    # these are the service account credentials returned from Martha
+    # these are the service account credentials returned from the DRS Resolver
     if sa_credentials:
         creds = service_account.Credentials.from_service_account_info(sa_credentials)
-    # if Martha did not give us a service account, try the credentials set at GOOGLE_APPLICATION_CREDENTIALS
+    # if the DRS Resolver did not give us a service account, try the credentials set at GOOGLE_APPLICATION_CREDENTIALS
     elif default_service_account_credentials:
         creds = service_account.Credentials.from_service_account_file(default_service_account_credentials)
     # we can't access the file

tests/test_drs.py

@@ -74,7 +74,7 @@ class TestTerraNotebookUtilsDRS(SuppressWarningsMixin, unittest.TestCase):
     jade_dev_url = "drs://jade.datarepo-dev.broadinstitute.org/v1_0c86170e-312d-4b39-a0a4-2a2bfaa24c7a_" \
                    "c0e40912-8b14-43f6-9a2f-b278144d0060"
 
-    # martha_v3 responses
+    # drs_resolver responses
     mock_jdr_response = {
         'contentType': 'application/octet-stream',
         'size': 15601108255,
@@ -92,7 +92,7 @@ class TestTerraNotebookUtilsDRS(SuppressWarningsMixin, unittest.TestCase):
             'crc32c': '8a366443'
         }
     }
-    mock_martha_v3_response_missing_fields = {
+    mock_drs_resolver_response_missing_fields = {
         'contentType': 'application/octet-stream',
         'bucket': 'broad-jade-dev-data-bucket',
         'name': 'fd8d8492-ad02-447d-b54e-35a7ffd0e7a5/8b07563a-542f-4b5c-9e00-e8fe6b1861de',
@@ -107,7 +107,7 @@ class TestTerraNotebookUtilsDRS(SuppressWarningsMixin, unittest.TestCase):
             'md5': '336ea55913bc261b72875bd259753046',
         }
     }
-    mock_martha_v3_response_without_gs_uri = {
+    mock_drs_resolver_response_without_gs_uri = {
         'contentType': 'application/octet-stream',
         'bucket': 'broad-jade-dev-data-bucket',
         'googleServiceAccount': {
@@ -119,7 +119,7 @@ class TestTerraNotebookUtilsDRS(SuppressWarningsMixin, unittest.TestCase):
             'md5': '336ea55913bc261b72875bd259753046',
         }
     }
-    mock_martha_v3_error_response = {
+    mock_drs_resolver_error_response = {
         "status": 500,
         "response": {
             "req": {
@@ -144,7 +144,7 @@ class TestTerraNotebookUtilsDRS(SuppressWarningsMixin, unittest.TestCase):
             "text": "{\"msg\":\"User 'null' does not have required action: read_data\",\"status_code\":500}"
         }
     }
-    mock_martha_v3_empty_error_response = {
+    mock_drs_resolver_empty_error_response = {
         "status": 500,
         "response": {
             "status": 500,
@@ -423,8 +423,8 @@ def test_arg_propagation(self):
                     drs.extract_tar_gz(self.drs_url)
                     enable_requester_pays.assert_called_with(WORKSPACE_NAME, WORKSPACE_NAMESPACE)
 
-    # test for when we get everything what we wanted in martha_v3 response
-    def test_martha_v3_response(self):
+    # test for when we get everything what we wanted in drs_resolver response
+    def test_drs_resolver_response(self):
         resp_json = mock.MagicMock(return_value=self.mock_jdr_response)
         requests_post = mock.MagicMock(return_value=mock.MagicMock(status_code=200, json=resp_json))
         with ExitStack() as es:
@@ -454,9 +454,9 @@ def test_martha_v2_response(self):
             self.assertEqual(15601108255, actual_info.size)
             self.assertEqual('2020-04-27T15:56:09.696Z', actual_info.updated)
 
-    # test for when some fields are missing in martha_v3 response
-    def test_martha_v3_response_with_missing_fields(self):
-        resp_json = mock.MagicMock(return_value=self.mock_martha_v3_response_missing_fields)
+    # test for when some fields are missing in drs_resolver response
+    def test_drs_resolver_response_with_missing_fields(self):
+        resp_json = mock.MagicMock(return_value=self.mock_drs_resolver_response_missing_fields)
         requests_post = mock.MagicMock(return_value=mock.MagicMock(status_code=200, json=resp_json))
         with ExitStack() as es:
             es.enter_context(mock.patch("terra_notebook_utils.drs.gs.get_client"))
@@ -485,10 +485,10 @@ def test_martha_v2_response_with_missing_fields(self):
             self.assertEqual(None, actual_info.size)
             self.assertEqual(None, actual_info.updated)
 
-    # test for when 'gsUrl' is missing in martha_v3 response. It should throw error
+    # test for when 'gsUrl' is missing in drs_resolver response. It should throw error
     @unittest.skip("TODO: Test no gsUri _and_ no accessUrl")
-    def test_martha_v3_response_without_gs_uri(self):
-        resp_json = mock.MagicMock(return_value=self.mock_martha_v3_response_without_gs_uri)
+    def test_drs_resolver_response_without_gs_uri(self):
+        resp_json = mock.MagicMock(return_value=self.mock_drs_resolver_response_without_gs_uri)
         requests_post = mock.MagicMock(return_value=mock.MagicMock(status_code=200, json=resp_json))
         with ExitStack() as es:
             es.enter_context(mock.patch("terra_notebook_utils.drs.gs.get_client"))
@@ -506,19 +506,19 @@ def test_martha_v2_response_without_gs_uri(self):
             with self.assertRaisesRegex(Exception, f"No GS url found for DRS uri '{self.drs_url}'"):
                 drs.get_drs_blob(self.drs_url)
 
-    # test for when martha_v3 returns error. It should throw error
-    def test_martha_v3_error_response(self):
-        resp_json = mock.MagicMock(return_value=self.mock_martha_v3_error_response)
+    # test for when drs_resolver returns error. It should throw error
+    def test_drs_resolver_error_response(self):
+        resp_json = mock.MagicMock(return_value=self.mock_drs_resolver_error_response)
         requests_post = mock.MagicMock(return_value=mock.MagicMock(status_code=500, json=resp_json))
         with ExitStack() as es:
             es.enter_context(mock.patch("terra_notebook_utils.drs.gs.get_client"))
             es.enter_context(mock.patch("terra_notebook_utils.drs.http", post=requests_post))
             with self.assertRaises(drs.DRSResolutionError):
                 drs.get_drs_blob(self.jade_dev_url)
 
-    # test for when martha_v3 returns error response with 'text' field. It should throw error
-    def test_martha_v3_empty_error_response(self):
-        resp_json = mock.MagicMock(return_value=self.mock_martha_v3_empty_error_response)
+    # test for when drs_resolver returns error response with 'text' field. It should throw error
+    def test_drs_resolver_empty_error_response(self):
+        resp_json = mock.MagicMock(return_value=self.mock_drs_resolver_empty_error_response)
         requests_post = mock.MagicMock(return_value=mock.MagicMock(status_code=500, json=resp_json))
         with ExitStack() as es:
             es.enter_context(mock.patch("terra_notebook_utils.drs.gs.get_client"))
@@ -570,7 +570,8 @@ def test_access(self):
             response = requests.get(signed_url, headers={'Range': 'bytes=0-1'})
             response.raise_for_status()
         # Test a Jade resource
-        # requires that GOOGLE_APPLICATION_CREDENTIALS be set, because martha does not return a service account
+        # requires that GOOGLE_APPLICATION_CREDENTIALS be set
+        # because neither martha nor DRSHub returns a service account
         jade_uri = 'drs://jade-terra.datarepo-prod.broadinstitute.org/' \
                    'v1_c3c588a8-be3f-467f-a244-da614be6889a_635984f0-3267-4201-b1ee-d82f64b8e6d1'
         with self.subTest(f'Testing DRS Access: {jade_uri}'):