Initial implementation of Azure auth support for DRS

Add support for using Microsoft Azure default credentials
when running in a Terra Azure Interactive Analysis Cloud Environment.
These credentials are used to auth with Terra backend services,
including Martha/terra-drs-hub and Rawls.

When running in a Terra Google IA Cloud Environment,
the behavior is as before, with no changes.

This initial implementation is thought to be functionally complete,
yet tests and documentation remain to be added.

Author: Michael Baumann

requirements.txt

@@ -1,3 +1,4 @@
+azure-identity >= 1.12.0, < 2
 google-cloud-storage >= 1.38.0, < 2
 gs-chunked-io >= 0.5.1, < 0.6
 firecloud

terra_notebook_utils/__init__.py

@@ -1,4 +1,6 @@
 import os
+from dataclasses import dataclass
+from enum import Enum
 
 WORKSPACE_NAME = os.environ.get('WORKSPACE_NAME', None)
 WORKSPACE_NAMESPACE = os.environ.get('WORKSPACE_NAMESPACE')  # This env var is set in Terra Cloud Environments
@@ -26,3 +28,20 @@
     DRS_RESOLVER_URL = MARTHA_URL
 else:
     DRS_RESOLVER_URL = DRSHUB_URL
+
+
+class ExecutionEnvironment(Enum):
+    TERRA_WORKSPACE = 1,  # Executing in a Terra Workspace (on any supported platform)
+    OTHER = 2  # Executing outside of a Terra Workspace (e.g., local system)
+
+
+class ExecutionPlatform(Enum):
+    AZURE = 1,  # Executing in an Azure compute environment
+    GOOGLE = 2,  # Executing in a Google compute environment
+    UNKNOWN = 3  # Execution platform not identified (e.g., local system)
+
+
+@dataclass
+class ExecutionContext:
+    execution_environment: ExecutionEnvironment
+    execution_platform: ExecutionPlatform

terra_notebook_utils/azure_auth.py

@@ -0,0 +1,58 @@
+"""
+Microsoft Azure specific code.
+Currently limited to Azure auth.
+
+See:
+https://azuresdkdocs.blob.core.windows.net/$web/python/azure-identity/1.12.0/index.html
+https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python
+https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/identity/azure-identity/azure/identity/_credentials/default.py
+"""
+
+import os
+from typing import Optional
+
+from azure.identity import DefaultAzureCredential
+from terra_notebook_utils.logger import logger
+
+
+# Single instance of DefaultAzureCredential that initialized lazily.
+# The instance is treated as threadsafe and reusable.
+# The Azure documentation is silent on thread safety.
+# Based on scanning the code, it appears to be threadsafe.
+# See: https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/identity/
+#              azure-identity/azure/identity/_credentials/default.py
+_AZURE_CREDENTIAL: Optional[DefaultAzureCredential] = None
+
+
+def _get_default_credential() -> DefaultAzureCredential:
+    """
+    Instantiate DefaultAzureCredential lazily if/when needed.
+
+    Note: It would not need to be instantiated this way, as
+    # no exception is raised even if Azure credentials are not configured.
+    :return: Reference to instance of DefaultAzureCredential
+    """
+
+    # Should a more sophisticated Singleton pattern be used instead?
+    global _AZURE_CREDENTIAL
+    if not _AZURE_CREDENTIAL:
+        _AZURE_CREDENTIAL = DefaultAzureCredential()
+    return _AZURE_CREDENTIAL
+
+
+def get_azure_access_token() -> str:
+    """
+    Return an Azure access token.
+
+    raises ClientAuthenticationError
+    """
+    if os.environ.get('TERRA_NOTEBOOK_AZURE_ACCESS_TOKEN'):
+        logger.debug("Using Azure token configured using 'TERRA_NOTEBOOK_AZURE_ACCESS_TOKEN'")
+        token = os.environ['TERRA_NOTEBOOK_AZURE_ACCESS_TOKEN']
+    else:
+        logger.debug("Requesting Azure default credentials token.")
+        token_scope = "https://management.azure.com/.default"
+        azure_token = _get_default_credential().get_token(token_scope)
+        logger.debug("Using Azure default credentials token.")
+        token = azure_token.token
+    return token

terra_notebook_utils/drs.py

@@ -11,13 +11,14 @@
     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
-from terra_notebook_utils.http import http
+from terra_notebook_utils.http_utils import http
 from terra_notebook_utils.blobstore.gs import GSBlob
 from terra_notebook_utils.blobstore.local import LocalBlob
 from terra_notebook_utils.blobstore.url import URLBlob
 from terra_notebook_utils.blobstore.progress import Indicator
 from terra_notebook_utils.blobstore import Blob, copy_client, BlobNotFoundError
 from terra_notebook_utils.logger import logger
+from terra_notebook_utils.terra_auth import get_terra_access_token
 
 
 DRSInfo = namedtuple("DRSInfo", "credentials access_url bucket_name key name size updated checksums")
@@ -46,7 +47,7 @@ def enable_requester_pays(workspace_name: Optional[str]=WORKSPACE_NAME,
     rawls_url = (f"https://rawls.dsde-{TERRA_DEPLOYMENT_ENV}.broadinstitute.org/api/workspaces/"
                  f"{workspace_namespace}/{encoded_workspace}/enableRequesterPaysForLinkedServiceAccounts")
     logger.info("Enabling requester pays for your workspace. This will only take a few seconds...")
-    access_token = gs.get_access_token()
+    access_token = get_terra_access_token()
 
     headers = {
         'authorization': f"Bearer {access_token}",
@@ -61,7 +62,7 @@ def enable_requester_pays(workspace_name: Optional[str]=WORKSPACE_NAME,
 
 def get_drs(drs_url: str, fields: List[str]) -> Response:
     """Request DRS information from DRS Resolver."""
-    access_token = gs.get_access_token()
+    access_token = get_terra_access_token()
 
     headers = {
         'authorization': f"Bearer {access_token}",

terra_notebook_utils/gs.py

@@ -42,7 +42,7 @@ def get_access_token():
     return token
 
 def reset_bond_cache():
-    from terra_notebook_utils.http import http
+    from terra_notebook_utils.http_utils import http
     token = get_access_token()
     headers = {
         'authorization': f'Bearer {token}',

terra_notebook_utils/http.py renamed to terra_notebook_utils/http_utils.py

@@ -1,4 +1,4 @@
-import json
 import logging
+from logging import Logger
 
-logger = logging.getLogger(__name__)
+logger: Logger = logging.getLogger(__name__)

terra_notebook_utils/logger.py

@@ -7,7 +7,7 @@
 
 import requests
 
-from terra_notebook_utils.http import Retry, http_session
+from terra_notebook_utils.http_utils import Retry, http_session
 from terra_notebook_utils.utils import _AsyncContextManager
 from terra_notebook_utils import WORKSPACE_NAME, WORKSPACE_NAMESPACE
 

terra_notebook_utils/table.py

@@ -0,0 +1,70 @@
+"""
+Support for auth with Terra backend services.
+"""
+from azure.core.exceptions import ClientAuthenticationError
+from google.auth.exceptions import DefaultCredentialsError
+
+from terra_notebook_utils import azure_auth, gs, ExecutionPlatform
+from terra_notebook_utils.logger import logger
+from terra_notebook_utils.utils import get_execution_context
+
+
+class AuthenticationError(Exception):
+    pass
+
+
+class TerraAuthTokenProvider:
+    """
+    Provides auth bearer tokens suitable for use with Terra backend services.
+    """
+    def __init__(self):
+        self.execution_context = get_execution_context()
+
+    @staticmethod
+    def _identify_valid_access_token() -> str:
+        """
+        Try to obtain an auth bearer tokens suitable for use with Terra backend services
+        from the Terra supported auth providers. First try Google, then try Azure.
+        Return the first successfully obtained token, otherwise raise AuthenticationError.
+
+        :return: auth bearer token suitable for use with Terra backend services
+        :raises: AuthenticationError
+        """
+        try:
+            logger.debug("Attempting to obtain a Google access token to use with Terra backend services.")
+            google_token = gs.get_access_token()
+            logger.debug("Using Google access token to use with Terra backend services.")
+            return google_token
+        except DefaultCredentialsError as ex:
+            logger.debug("Failed to obtain a Google access token to use with Terra backend services.", exc_info=ex)
+
+        try:
+            logger.debug("Attempting to obtain a Azure access token to use with Terra backend services.")
+            azure_token = azure_auth.get_azure_access_token()
+            logger.debug("Using Azure access token to use with Terra backend services.")
+            return azure_token
+        except ClientAuthenticationError as ex:
+            logger.debug("Failed to obtain a Azure access token to use with Terra backend services.", exc_info=ex)
+
+        raise AuthenticationError("Failed to obtain a Google or Azure token to auth with Terra backend services.")
+
+    def get_terra_access_token(self) -> str:
+        if self.execution_context.execution_platform == ExecutionPlatform.GOOGLE:
+            logger.debug("Using Google default credentials to auth with Terra services.")
+            return gs.get_access_token()
+        elif self.execution_context.execution_platform == ExecutionPlatform.AZURE:
+            logger.debug("Using Azure default credentials to auth with Terra services.")
+            return azure_auth.get_azure_access_token()
+        else:
+            return self._identify_valid_access_token()
+
+
+# Single instance of TerraAuthTokenProvider.
+TERRA_AUTH_TOKEN_PROVIDER = TerraAuthTokenProvider()
+
+
+def get_terra_access_token() -> str:
+    """ Return an auth bearer token suitable for use with Terra backend services.
+    :raises: AuthenticationError
+    """
+    return TERRA_AUTH_TOKEN_PROVIDER.get_terra_access_token()

terra_notebook_utils/terra_auth.py

@@ -1,11 +1,14 @@
 import json
+import os
 import threading
 from functools import lru_cache
 from concurrent.futures import ThreadPoolExecutor, Future, as_completed
 from typing import Any, Callable, Dict, Optional, Iterable, Set
 
 import jmespath
 
+from terra_notebook_utils import ExecutionEnvironment, ExecutionPlatform, ExecutionContext
+
 
 class _AsyncContextManager:
     """Context manager for asynchronous execution. Wait on exit for all jobs to complete."""
@@ -63,3 +66,22 @@ def is_notebook() -> bool:
         return "ZMQInteractiveShell" == get_ipython().__class__.__name__  # type: ignore
     except NameError:
         return False
+
+
+@lru_cache()
+def get_execution_context() -> ExecutionContext:
+    """
+    Identify information about the context in which terra-notebook-utils is executing.
+    Currently, this determination is made based on the presence and value of the WORKSPACE_BUCKET.
+    Other/improved means may be identified in the future.
+    """
+    execution_environment = ExecutionEnvironment.OTHER
+    execution_platform = ExecutionPlatform.UNKNOWN
+    workspace_bucket = os.environ.get('WORKSPACE_BUCKET', None)
+    if workspace_bucket:
+        execution_environment = ExecutionEnvironment.TERRA_WORKSPACE
+        if workspace_bucket.startswith("gs://"):
+            execution_platform = ExecutionPlatform.GOOGLE
+        elif workspace_bucket.startswith("https://"):
+            execution_platform = ExecutionPlatform.AZURE
+    return ExecutionContext(execution_environment, execution_platform)