Merge pull request #186 from statisticsnorway/add-dapla-lab-email

Fix 'fetch_email_from_credentials' for Dapla Lab

Author: mallport

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dapla-toolbelt"
-version = "3.1.2"
+version = "3.2.0"
 description = "Dapla Toolbelt"
 authors = ["Dapla Developers <dapla-platform-developers@ssb.no>"]
 license = "MIT"

src/dapla/auth.py

@@ -5,7 +5,6 @@
 from collections.abc import Sequence
 from datetime import datetime
 from datetime import timedelta
-from enum import Enum
 from functools import lru_cache
 from functools import partial
 from typing import Any
@@ -19,6 +18,10 @@
 from IPython.display import display
 from jupyterhub.services.auth import HubAuth
 
+from dapla.const import DaplaEnvironment
+from dapla.const import DaplaRegion
+from dapla.const import DaplaService
+
 logger = logging.getLogger(__name__)
 
 # Refresh window was modified in: https://github.com/googleapis/google-auth-library-python/commit/c6af1d692b43833baca978948376739547cf685a
@@ -29,33 +32,6 @@
 google.auth._helpers.REFRESH_THRESHOLD = timedelta(seconds=20)
 
 
-class DaplaEnvironment(Enum):
-    """Represents the 'DAPLA_ENVIRONMENT' environment variable."""
-
-    DEV = "DEV"
-    STAGING = "STAGING"
-    TEST = "TEST"
-    PROD = "PROD"
-
-
-class DaplaService(Enum):
-    """Represents the 'DAPLA_SERVICE' environment variable."""
-
-    JUPYTERLAB = "JUPYTERLAB"
-    VS_CODE = "VS_CODE"
-    R_STUDIO = "R_STUDIO"
-    CLOUD_RUN = "CLOUD_RUN"
-
-
-class DaplaRegion(Enum):
-    """Represents the 'DAPLA_REGION' environment variable."""
-
-    ON_PREM = "ON_PREM"
-    DAPLA_LAB = "DAPLA_LAB"
-    BIP = "BIP"
-    CLOUD_RUN = "CLOUD_RUN"
-
-
 class AuthClient:
     """Client for retrieving authentication information."""
 
@@ -347,6 +323,9 @@ def fetch_personal_token() -> str:
     @lru_cache(maxsize=1)
     def fetch_email_from_credentials() -> Optional[str]:
         """Retrieves an e-mail based on current Google Credentials. Potentially makes a Google API call."""
+        if os.getenv("DAPLA_REGION") == str(DaplaRegion.DAPLA_LAB):
+            return os.getenv("DAPLA_USER")
+
         credentials = AuthClient.fetch_google_credentials()
         response = requests.get(
             url=f"https://oauth2.googleapis.com/tokeninfo?access_token={credentials.token}"

src/dapla/const.py

@@ -0,0 +1,28 @@
+from enum import Enum
+
+
+class DaplaEnvironment(Enum):
+    """Represents the 'DAPLA_ENVIRONMENT' environment variable."""
+
+    DEV = "DEV"
+    STAGING = "STAGING"
+    TEST = "TEST"
+    PROD = "PROD"
+
+
+class DaplaService(Enum):
+    """Represents the 'DAPLA_SERVICE' environment variable."""
+
+    JUPYTERLAB = "JUPYTERLAB"
+    VS_CODE = "VS_CODE"
+    R_STUDIO = "R_STUDIO"
+    CLOUD_RUN = "CLOUD_RUN"
+
+
+class DaplaRegion(Enum):
+    """Represents the 'DAPLA_REGION' environment variable."""
+
+    ON_PREM = "ON_PREM"
+    DAPLA_LAB = "DAPLA_LAB"
+    BIP = "BIP"
+    CLOUD_RUN = "CLOUD_RUN"

src/dapla/doctor.py

@@ -8,7 +8,7 @@
 from gcsfs.retry import HttpError
 
 from dapla.auth import AuthClient
-from dapla.auth import DaplaRegion
+from dapla.const import DaplaRegion
 
 logger = logging.getLogger(__name__)
 

tests/test_pubsub.py

@@ -6,7 +6,7 @@
 import pytest
 from google.api_core.exceptions import NotFound
 from google.auth.exceptions import DefaultCredentialsError
-from google.cloud import pubsub_v1
+from google.cloud.pubsub_v1.publisher.futures import Future
 
 import dapla.pubsub
 from dapla.pubsub import _extract_env
@@ -55,9 +55,7 @@ def test_publish_gcs_objects_to_pubsub(self, mock_list: Mock) -> None:
             )
 
     def test_get_callback(self) -> None:
-        publish_future = MagicMock(
-            side_effect=pubsub_v1.publisher.futures.Future.result
-        )
+        publish_future = MagicMock(side_effect=Future.result)
         # Create a callback function using the _get_callback helper function
         callback = _get_callback(publish_future, "blob_name", timeout=1)