Add support for CI use for feature flags

Add support for CI use for feature flags

Author: KevinFairise2

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

@@ -22,7 +22,12 @@ class AppEnvVars:
     VERBOSE = "DDA_VERBOSE"
     NO_DYNAMIC_DEPS = "DDA_NO_DYNAMIC_DEPS"
     TELEMETRY_API_KEY = "DDA_TELEMETRY_API_KEY"
-    FEATURE_FLAGS_CLIENT_TOKEN = "DDA_FEATURE_FLAGS_CLIENT_TOKEN"  # noqa: S105i
+    FEATURE_FLAGS_CLIENT_TOKEN = "DDA_FEATURE_FLAGS_CLIENT_TOKEN"  # noqa: S105
+    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,10 +26,9 @@ 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:
@@ -34,11 +39,129 @@ def __init__(self, app: Application) -> None:
         # 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] = {}
 
+    def enabled(self, flag: str, *, default: bool = False, scopes: Optional[dict[str, str]] = None) -> bool:
+        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 targeting key {entity} and targeting attributes {tuple_attributes}"
+        )
+        flag_value = self.__check_flag(flag, entity, tuple_attributes)
+        if flag_value is None:
+            return default
+        return flag_value
+
+    def __check_flag(self, flag: str, entity: str, scopes: tuple[tuple[str, str], ...]) -> bool | None:
+        if self.__client is None:
+            self.__app.display_debug("Feature flag client not initialized properly")
+            return None
+
+        cache_key = (flag, entity, scopes)
+        if cache_key in self.__cache:
+            return self.__cache[cache_key]
+
+        flag_value = self.__client.get_flag_value(flag, entity, dict(scopes))
+        self.__cache[cache_key] = flag_value
+        return flag_value
+
+    @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
+
+    @abstractmethod
+    def _get_client_token(self) -> str | None:
+        pass
+
     @cached_property
     def __client(self) -> DatadogFeatureFlag | None:
-        if running_in_ci():  # We do not support feature flags token retrieval in the CI yet.
+        token = self._get_client_token()
+        if token is None:
             return None
+        return DatadogFeatureFlag(token, self.__app)
 
+
+    @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
@@ -54,16 +177,12 @@ def __client(self) -> DatadogFeatureFlag | None:
 
     @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 FeatureFlagUser(self._app.config)
 
+    def _get_entity(self) -> str:
         return self.__user.machine_id
 
+<<<<<<< HEAD
     def enabled(self, flag: str, *, default: bool = False, scopes: Optional[dict[str, str]] = None) -> bool:
         entity = self.__get_entity()
         base_scopes = self.__get_base_scopes()
@@ -93,7 +212,11 @@ def __check_flag(self, flag: str, entity: str, scopes: tuple[tuple[str, str], ..
         return flag_value
 
     def __get_base_scopes(self) -> dict[str, str]:
+=======
+    def _get_base_scopes(self) -> dict[str, str]:
+>>>>>>> be7a4ee (Add support for CI use for feature flags)
         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"]