Support feature flags in the CI

pyproject.toml

@@ -26,6 +26,7 @@ classifiers = [
 dependencies = [
   "ada-url~=1.15.3",
   "binary~=1.0.2",
+  "boto3~=1.38.8",
   "click~=8.1",
   "datadog-api-client~=2.34",
   "dep-sync~=0.1",

src/dda/cli/application.py

@@ -9,6 +9,8 @@
 
 from dda.cli.terminal import Terminal
 from dda.config.constants import AppEnvVars
+from dda.feature_flags.manager import CIFeatureFlagManager, LocalFeatureFlagManager
+from dda.utils.ci import running_in_ci
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -129,9 +131,10 @@ def telemetry(self) -> TelemetryManager:
 
     @cached_property
     def features(self) -> FeatureFlagManager:
-        from dda.feature_flags.manager import FeatureFlagManager
-
-        return FeatureFlagManager(self)
+        self.display_info(f"IS IT RUNNING IN CI {running_in_ci()}")
+        if running_in_ci():
+            return CIFeatureFlagManager(self)
+        return LocalFeatureFlagManager(self)
 
     @cached_property
     def dynamic_deps_allowed(self) -> bool:

src/dda/config/constants.py

@@ -23,6 +23,11 @@ class AppEnvVars:
     NO_DYNAMIC_DEPS = "DDA_NO_DYNAMIC_DEPS"
     TELEMETRY_API_KEY = "DDA_TELEMETRY_API_KEY"
     FEATURE_FLAGS_CLIENT_TOKEN = "DDA_FEATURE_FLAGS_CLIENT_TOKEN"  # noqa: S105 This is not a hardcoded secret but the linter complains on it
+    FEATURE_FLAGS_CI_VAULT_PATH = "DDA_FEATURE_FLAGS_CI_VAULT_PATH"
+    FEATURE_FLAGS_CI_VAULT_KEY = "DDA_FEATURE_FLAGS_CI_VAULT_KEY"
+    FEATURE_FLAGS_CI_SSM_KEY_WINDOWS = "DDA_FEATURE_FLAGS_CI_SSM_KEY_WINDOWS"
+    FEATURE_FLAGS_CI_VAULT_KEY_MACOS = "DDA_FEATURE_FLAGS_CI_VAULT_KEY_MACOS"
+    FEATURE_FLAGS_CI_VAULT_PATH_MACOS = "DDA_FEATURE_FLAGS_CI_VAULT_PATH_MACOS"
     TELEMETRY_USER_MACHINE_ID = "DDA_TELEMETRY_USER_MACHINE_ID"
     # https://no-color.org
     NO_COLOR = "NO_COLOR"

src/dda/feature_flags/manager.py

@@ -3,12 +3,18 @@
 # SPDX-License-Identifier: MIT
 from __future__ import annotations
 
+import os
+import sys
+from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Optional
 
+from dda.config.constants import AppEnvVars
 from dda.feature_flags.client import DatadogFeatureFlag
+from dda.secrets.ssm import fetch_secret_ssm
+from dda.secrets.vault import fetch_secret_ci
 from dda.user.datadog import User
-from dda.utils.ci import running_in_ci
+from dda.utils.platform import get_os_name
 
 if TYPE_CHECKING:
     from dda.cli.application import Application
@@ -20,67 +26,38 @@ def __init__(self, config: RootConfig) -> None:
         super().__init__(config)
 
 
-class FeatureFlagManager:
+class FeatureFlagManager(ABC):
     """
-    A class for querying feature flags. This is available as the
-    [`Application.features`][dda.cli.application.Application.features] property.
+    A class for querying feature flags.
     """
 
     def __init__(self, app: Application) -> None:
-        self.__app = app
+        self._app = app
 
-        self.__client = DatadogFeatureFlag(self.__client_token, self.__app)
+        self.__client = DatadogFeatureFlag(self.__client_token, self._app)
 
         # Manually implemented cache to avoid calling several time Feature flag backend on the same flag evaluation.
         # Cache key is a tuple of the flag, entity and scopes, to make it hashable.
         # For example after calling `enabled("test-flag", default=False, scopes={"user": "user1"}),
         # the cache will contain the result for the tuple ("test-flag", "entity", (("user", "user1"),)).
         self.__cache: dict[tuple[str, str, tuple[tuple[str, str], ...]], Any] = {}
 
-    @cached_property
-    def __client_token(self) -> str | None:
-        if running_in_ci():  # We do not support feature flags token retrieval in the CI yet.
-            return None
-
-        from contextlib import suppress
-
-        from dda.secrets.api import fetch_client_token, read_client_token, save_client_token
-
-        client_token: str | None = None
-        with suppress(Exception):
-            client_token = read_client_token()
-            if not client_token:
-                client_token = fetch_client_token()
-                save_client_token(client_token)
-
-        return client_token
-
-    @property
-    def __user(self) -> FeatureFlagUser:
-        return FeatureFlagUser(self.__app.config)
-
-    def __get_entity(self) -> str:
-        if running_in_ci():
-            import os
-
-            return os.getenv("CI_JOB_ID", "default_job_id")
-
-        return self.__user.machine_id
-
     def enabled(self, flag: str, *, default: bool = False, scopes: Optional[dict[str, str]] = None) -> bool:
         if not self.__client_token:
-            self.__app.display_debug("No client token found")
+            self._app.display_debug("No client token found")
             return default
 
-        entity = self.__get_entity()
-        base_scopes = self.__get_base_scopes()
+        entity = self.__entity
+        base_scopes = self.__base_scopes
         if scopes is not None:
             base_scopes.update(scopes)
 
         attributes_items = base_scopes.items()
         tuple_attributes = tuple(((key, value) for key, value in sorted(attributes_items)))
 
-        self.__app.display_debug(f"Checking flag {flag} with entity {entity} and scopes {base_scopes}")
+        self._app.display_debug(
+            f"Checking flag {flag} with targeting key {entity} and targeting attributes {tuple_attributes}"
+        )
         flag_value = self.__check_flag(flag, entity, tuple_attributes)
         if flag_value is None:
             return default
@@ -95,8 +72,114 @@ def __check_flag(self, flag: str, entity: str, scopes: tuple[tuple[str, str], ..
         self.__cache[cache_key] = flag_value
         return flag_value
 
-    def __get_base_scopes(self) -> dict[str, str]:
+    @abstractmethod
+    def _get_client_token(self) -> str | None:
+        pass
+
+    @abstractmethod
+    def _get_entity(self) -> str:
+        pass
+
+    @abstractmethod
+    def _get_base_scopes(self) -> dict[str, str]:
+        pass
+
+    @cached_property
+    def __client_token(self) -> str | None:
+        return self._get_client_token()
+
+    @cached_property
+    def __base_scopes(self) -> dict[str, str]:
+        return self._get_base_scopes()
+
+    @cached_property
+    def __entity(self) -> str:
+        return self._get_entity()
+
+
+class CIFeatureFlagManager(FeatureFlagManager):
+    """
+    A class for querying feature flags in a CI environment.
+    """
+
+    def _get_client_token(self) -> str | None:
+        self._app.display_debug(f"Getting client token for {sys.platform}")
+        try:
+            match sys.platform:
+                case "win32":
+                    return self.__get_client_token_windows()
+                case "darwin":
+                    return self.__get_client_token_macos()
+                case "linux":
+                    return self.__get_client_token_linux()
+                case _:
+                    return None
+        except Exception as e:  # noqa: BLE001
+            self._app.display_warning(f"Error getting client token: {e}, feature flag will be defaulted")
+            return None
+
+    def __get_client_token_windows(self) -> str | None:  # noqa: PLR6301
+        if (client_token := os.getenv(AppEnvVars.FEATURE_FLAGS_CI_SSM_KEY_WINDOWS)) is None:
+            return None
+        return fetch_secret_ssm(name=client_token)
+
+    def __get_client_token_macos(self) -> str | None:  # noqa: PLR6301
+        if (client_token := os.getenv(AppEnvVars.FEATURE_FLAGS_CI_VAULT_KEY_MACOS)) is None:
+            return None
+        if (vault_path := os.getenv(AppEnvVars.FEATURE_FLAGS_CI_VAULT_PATH_MACOS)) is None:
+            return None
+        return fetch_secret_ci(vault_path, client_token)
+
+    def __get_client_token_linux(self) -> str | None:  # noqa: PLR6301
+        if (client_token := os.getenv(AppEnvVars.FEATURE_FLAGS_CI_VAULT_KEY)) is None:
+            return None
+        if (vault_path := os.getenv(AppEnvVars.FEATURE_FLAGS_CI_VAULT_PATH)) is None:
+            return None
+        return fetch_secret_ci(vault_path, client_token)
+
+    def _get_entity(self) -> str:  # noqa: PLR6301
+        return os.getenv("CI_JOB_ID", "default_entity")
+
+    def _get_base_scopes(self) -> dict[str, str]:  # noqa: PLR6301
+        return {
+            "ci": "true",
+            "ci.job.name": os.getenv("CI_JOB_NAME", ""),
+            "ci.job.id": os.getenv("CI_JOB_ID", ""),
+            "ci.stage.name": os.getenv("CI_JOB_STAGE", ""),
+            "git.branch": os.getenv("CI_COMMIT_BRANCH", ""),
+        }
+
+
+class LocalFeatureFlagManager(FeatureFlagManager):
+    """
+    A class for querying feature flags. This is available as the
+    [`Application.features`][dda.cli.application.Application.features] property.
+    """
+
+    def _get_client_token(self) -> str | None:  # noqa: PLR6301
+        from contextlib import suppress
+
+        from dda.secrets.api import fetch_client_token, read_client_token, save_client_token
+
+        client_token: str | None = None
+        with suppress(Exception):
+            client_token = read_client_token()
+            if not client_token:
+                client_token = fetch_client_token()
+                save_client_token(client_token)
+
+        return client_token
+
+    @property
+    def __user(self) -> FeatureFlagUser:
+        return FeatureFlagUser(self._app.config)
+
+    def _get_entity(self) -> str:
+        return self.__user.machine_id
+
+    def _get_base_scopes(self) -> dict[str, str]:
         return {
-            "ci": "true" if running_in_ci() else "false",
+            "platform": get_os_name(),
+            "ci": "false",
             "user": self.__user.email,
         }

src/dda/secrets/ssm.py

@@ -0,0 +1,6 @@
+def fetch_secret_ssm(name: str) -> str:
+    import boto3
+
+    ssm = boto3.client("ssm", region_name="us-east-1")
+    response = ssm.get_parameter(Name=name, WithDecryption=True)
+    return response["Parameter"]["Value"]

src/dda/secrets/vault.py

@@ -10,6 +10,8 @@
 import hvac
 from ada_url import URL, URLSearchParams
 
+from dda.utils.ci import running_in_ci
+
 VAULT_URL = "https://vault.us1.ddbuild.io"
 OIDC_CALLBACK_PORT = 8250
 OIDC_REDIRECT_URI = f"http://localhost:{OIDC_CALLBACK_PORT}/oidc/callback"
@@ -94,6 +96,18 @@ def do_GET(self) -> None:  # noqa: N802
 
 
 def fetch_secret(name: str, key: str) -> str:
+    if running_in_ci():
+        return fetch_secret_ci(name, key)
+    return fetch_secret_local(name, key)
+
+
+def fetch_secret_local(name: str, key: str) -> str:
     client = init_client()
     secret = client.secrets.kv.v2.read_secret_version(path=name, mount_point="kv")
     return secret["data"]["data"][key]
+
+
+def fetch_secret_ci(name: str, key: str) -> str:
+    client = hvac.Client()
+    secret = client.secrets.kv.v2.read_secret_version(path=name, mount_point="kv")
+    return secret["data"]["data"][key]

src/dda/telemetry/daemon/trace.py

@@ -15,37 +15,13 @@
 from dda.telemetry.constants import SERVICE_NAME, SERVICE_VERSION
 from dda.telemetry.daemon.base import TelemetryClient
 from dda.utils.network.http.client import get_http_client
-from dda.utils.platform import get_machine_id
+from dda.utils.platform import get_machine_id, get_os_name, get_os_version
 
 if TYPE_CHECKING:
     from types import TracebackType
 
 URL = "https://instrumentation-telemetry-intake.datadoghq.com/api/v2/apmtelemetry"
 
-if sys.platform == "win32":
-
-    def get_os_name() -> str:
-        return f"{platform.system()} {platform.win32_ver()[0]} {platform.win32_edition()}"
-
-    def get_os_version() -> str:
-        return platform.win32_ver()[0]
-
-elif sys.platform == "darwin":
-
-    def get_os_name() -> str:
-        return f"{platform.system()} {platform.mac_ver()[0]}"
-
-    def get_os_version() -> str:
-        return platform.mac_ver()[0]
-
-else:
-
-    def get_os_name() -> str:
-        return platform.freedesktop_os_release()["NAME"]
-
-    def get_os_version() -> str:
-        return platform.freedesktop_os_release()["VERSION_ID"]
-
 
 @cache
 def get_base_payload() -> dict[str, Any]:

src/dda/utils/platform/__init__.py

@@ -5,6 +5,7 @@
 
 import contextlib
 import os
+import platform
 import sys
 from functools import cache
 from typing import TYPE_CHECKING
@@ -174,3 +175,28 @@ def get_machine_id() -> UUID:
     import uuid
 
     return uuid.uuid5(uuid.NAMESPACE_DNS, str(uuid.getnode()))
+
+
+if sys.platform == "win32":
+
+    def get_os_name() -> str:
+        return f"{platform.system()} {platform.win32_ver()[0]} {platform.win32_edition()}"
+
+    def get_os_version() -> str:
+        return platform.win32_ver()[0]
+
+elif sys.platform == "darwin":
+
+    def get_os_name() -> str:
+        return f"{platform.system()} {platform.mac_ver()[0]}"
+
+    def get_os_version() -> str:
+        return platform.mac_ver()[0]
+
+else:
+
+    def get_os_name() -> str:
+        return platform.freedesktop_os_release()["NAME"]
+
+    def get_os_version() -> str:
+        return platform.freedesktop_os_release()["VERSION_ID"]