Browse Source

ref(flagpole): Copies sentry flagpole context into sentry src directory (#72059)

Gabe Villalobos 9 months ago
parent
commit
ace8aeba85

+ 6 - 2
src/flagpole/evaluation_context.py

@@ -64,8 +64,12 @@ class ContextBuilder(Generic[T_CONTEXT_DATA]):
     >>> feature.match(dict())
     """
 
-    context_transformers: list[Callable[[T_CONTEXT_DATA], EvaluationContextDict]] = []
-    exception_handler: Callable[[Exception], Any] | None = None
+    context_transformers: list[Callable[[T_CONTEXT_DATA], EvaluationContextDict]]
+    exception_handler: Callable[[Exception], Any] | None
+
+    def __init__(self):
+        self.context_transformers = []
+        self.exception_handler = None
 
     def add_context_transformer(
         self, context_transformer: Callable[[T_CONTEXT_DATA], EvaluationContextDict]

+ 102 - 0
src/sentry/features/flagpole_context.py

@@ -0,0 +1,102 @@
+from dataclasses import dataclass
+
+from django.contrib.auth.models import AnonymousUser
+
+from flagpole.evaluation_context import ContextBuilder, EvaluationContextDict
+from sentry.models.organization import Organization
+from sentry.models.project import Project
+from sentry.models.user import User
+from sentry.services.hybrid_cloud.organization import RpcOrganization
+from sentry.services.hybrid_cloud.project import RpcProject
+from sentry.services.hybrid_cloud.user import RpcUser
+
+
+class InvalidContextDataException(Exception):
+    pass
+
+
+@dataclass()
+class SentryContextData:
+    actor: User | RpcUser | AnonymousUser | None = None
+    organization: Organization | RpcOrganization | None = None
+    project: Project | RpcProject | None = None
+
+
+def organization_context_transformer(data: SentryContextData) -> EvaluationContextDict:
+    context_data: EvaluationContextDict = dict()
+    org = data.organization
+    if org is None:
+        return context_data
+
+    if isinstance(org, Organization):
+        context_data["organization_slug"] = org.slug
+        context_data["organization_name"] = org.name
+        context_data["organization_id"] = org.id
+        early_adopter = bool(org.flags.early_adopter) if org.flags is not None else False
+        context_data["organization_is-early-adopter"] = early_adopter
+
+    elif isinstance(org, RpcOrganization):
+        context_data["organization_slug"] = org.slug
+        context_data["organization_name"] = org.name
+        context_data["organization_id"] = org.id
+        context_data["organization_is-early-adopter"] = org.flags.early_adopter
+    else:
+        raise InvalidContextDataException("Invalid organization object provided")
+
+    return context_data
+
+
+def project_context_transformer(data: SentryContextData) -> EvaluationContextDict:
+    context_data: EvaluationContextDict = dict()
+
+    if (proj := data.project) is not None:
+        if not isinstance(proj, Project):
+            raise InvalidContextDataException("Invalid project object provided")
+
+        context_data["project_slug"] = proj.slug
+        context_data["project_name"] = proj.name
+        context_data["project_id"] = proj.id
+
+    return context_data
+
+
+def user_context_transformer(data: SentryContextData) -> EvaluationContextDict:
+    context_data: EvaluationContextDict = dict()
+    user = data.actor
+    if user is None or isinstance(user, AnonymousUser):
+        return context_data
+
+    if not isinstance(user, User) and not isinstance(user, RpcUser):
+        raise InvalidContextDataException("Invalid actor object provided")
+
+    if user.is_authenticated:
+        context_data["user_id"] = user.id
+        context_data["user_is-superuser"] = user.is_superuser
+        context_data["user_is-staff"] = user.is_staff
+
+    verified_emails: list[str]
+
+    if isinstance(user, RpcUser):
+        verified_emails = list(user.emails)
+    else:
+        verified_emails = user.get_verified_emails().values_list("email", flat=True)
+
+    if user.email in verified_emails:
+        context_data["user_email"] = user.email
+        context_data["user_domain"] = user.email.rsplit("@", 1)[-1]
+
+    return context_data
+
+
+def get_sentry_flagpole_context_builder() -> ContextBuilder[SentryContextData]:
+    """
+    Creates and returns a new sentry flagpole context builder with Organization,
+     User, and Project transformers appended to it.
+    :return:
+    """
+    return (
+        ContextBuilder[SentryContextData]()
+        .add_context_transformer(organization_context_transformer)
+        .add_context_transformer(project_context_transformer)
+        .add_context_transformer(user_context_transformer)
+    )

+ 3 - 3
tests/flagpole/test_feature.py

@@ -8,7 +8,7 @@ from flagpole.conditions import ConditionOperatorKind
 
 
 @dataclass
-class ContextData:
+class SimpleTestContextData:
     pass
 
 
@@ -110,7 +110,7 @@ class TestParseFeatureConfig:
         )
 
         context_builder = self.get_is_true_context_builder(is_true_value=True)
-        assert feature.match(context_builder.build(ContextData()))
+        assert feature.match(context_builder.build(SimpleTestContextData()))
 
     def test_disabled_feature(self):
         feature = Feature.from_feature_config_json(
@@ -134,4 +134,4 @@ class TestParseFeatureConfig:
         )
 
         context_builder = self.get_is_true_context_builder(is_true_value=True)
-        assert not feature.match(context_builder.build(ContextData()))
+        assert not feature.match(context_builder.build(SimpleTestContextData()))

+ 1 - 1
tests/flagpole/test_sentry_flagpole_context.py → tests/sentry/features/test_flagpole_context.py

@@ -1,7 +1,7 @@
 import pytest
 from django.contrib.auth.models import AnonymousUser
 
-from flagpole.sentry_flagpole_context import (
+from sentry.features.flagpole_context import (
     InvalidContextDataException,
     SentryContextData,
     get_sentry_flagpole_context_builder,