feat(seer-infra-telemetry): Add DatadogIdentityProvider for OAuth2

src/sentry/identity/__init__.py

@@ -11,6 +11,7 @@
 
 def _register_providers() -> None:
     from .bitbucket.provider import BitbucketIdentityProvider
+    from .datadog.provider import DatadogIdentityProvider
     from .discord.provider import DiscordIdentityProvider
     from .gcp.provider import GCPIdentityProvider
     from .github.provider import GitHubIdentityProvider
@@ -34,6 +35,7 @@ def _register_providers() -> None:
     register(GitlabIdentityProvider)
     register(GoogleIdentityProvider)
     register(DiscordIdentityProvider)
+    register(DatadogIdentityProvider)
     register(GCPIdentityProvider)
 
 

src/sentry/identity/datadog/provider.py

@@ -2,30 +2,55 @@
 
 import base64
 import hashlib
+import logging
 import secrets
+from typing import Any
 
 import orjson
 from django.http.request import HttpRequest
 from django.http.response import HttpResponseBase
 from requests import ConnectionError, HTTPError, Response
 from requests.exceptions import SSLError
 
+from sentry.auth.exceptions import IdentityNotValid
 from sentry.http import safe_urlopen, safe_urlread
 from sentry.identity.oauth2 import (
     OAuth2CallbackView,
     OAuth2LoginView,
+    OAuth2Provider,
     _redirect_url,
     record_event,
 )
 from sentry.identity.pipeline import IdentityPipeline
+from sentry.identity.services.identity.model import RpcIdentity
+from sentry.integrations.types import IntegrationProviderSlug
 from sentry.integrations.utils.metrics import IntegrationPipelineViewType
+from sentry.pipeline.views.base import PipelineView
+from sentry.shared_integrations.exceptions import ApiError, ApiInvalidRequestError, ApiUnauthorized
+from sentry.users.models.identity import Identity
 from sentry.utils.http import absolute_uri
 
+logger = logging.getLogger(__name__)
+
+MCP_REGISTER_PATH = "/api/unstable/mcp-server/register"
+MCP_AUTHORIZE_PATH = "/api/unstable/mcp-server/authorize"
+MCP_TOKEN_PATH = "/api/unstable/mcp-server/token"
+
 
 def _basic_auth_header(client_id: str, client_secret: str) -> str:
     return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode("ascii")
 
 
+def get_user_info(access_token: str, site: str) -> dict[str, Any]:
+    """Fetch the current Datadog user via ``GET /api/v2/current_user``."""
+    url = f"https://api.{site}/api/v2/current_user"
+    resp = safe_urlopen(url, method="GET", headers={"Authorization": f"Bearer {access_token}"})
+    resp.raise_for_status()
+
+    body = orjson.loads(safe_urlread(resp))
+    return body["data"]
+
+
 def generate_pkce_code_verifier() -> str:
     return secrets.token_urlsafe(96)
 
@@ -180,3 +205,108 @@
         return safe_urlopen(
             self.access_token_url, data=data, headers=headers, verify_ssl=verify_ssl
         )
+
+
+class DatadogIdentityProvider(OAuth2Provider):
+    key = IntegrationProviderSlug.DATADOG
+    name = "Datadog"
+
+    oauth_scopes: tuple[str, ...] = (
+        "apm_read",
+        "error_tracking_read",
+        "events_read",
+        "hosts_read",
+        "incident_read",
+        "logs_read_data",
+        "metrics_read",
+        "monitors_read",
+        "user_self_profile_read",
+    )
+
+    def _get_mcp_base_url(self) -> str:
+        return f"https://mcp.{self._get_oauth_parameter('site')}"
+
+    def get_oauth_authorize_url(self) -> str:
+        return self._get_mcp_base_url() + MCP_AUTHORIZE_PATH
+
+    def get_oauth_access_token_url(self) -> str:
+        return self._get_mcp_base_url() + MCP_TOKEN_PATH
+
+    def get_pipeline_views(self) -> list[PipelineView[IdentityPipeline]]:
+        return [
+            DatadogOAuth2DCRView(
+                register_url=self._get_mcp_base_url() + MCP_REGISTER_PATH,
+            ),
+            DatadogOAuth2LoginView(
+                authorize_url=self.get_oauth_authorize_url(),
+                scope=",".join(self.get_oauth_scopes()),
+                resource=self._get_mcp_base_url(),
+            ),
+            DatadogOAuth2CallbackView(
+                access_token_url=self.get_oauth_access_token_url(),
+            ),
+        ]
+
+    def get_oauth_data(self, payload: dict[str, Any]) -> dict[str, Any]:
+        data = super().get_oauth_data(payload)
+        if "scope" in payload:
+            data["scope"] = payload["scope"]
+        return data
+
+    def build_identity(self, data: dict[str, Any]) -> dict[str, Any]:
+        token_data = data["data"]
+        access_token = token_data.get("access_token")
+        if not access_token:
+            raise ValueError("Datadog token exchange did not return an access_token")
+
+        user = get_user_info(access_token, self._get_oauth_parameter("site"))
+
+        oauth_data = self.get_oauth_data(token_data)
+        # Persist DCR credentials so refresh_identity can authenticate later.
+        oauth_data["client_id"] = data.get("dcr_client_id")
+        oauth_data["client_secret"] = data.get("dcr_client_secret")
+
+        return {
+            "type": IntegrationProviderSlug.DATADOG,
+            "id": user["id"],
+            "email": user.get("attributes", {}).get("email"),
+            "name": user.get("attributes", {}).get("name"),
+            "scopes": [],
+            "data": oauth_data,
+        }
+
+    def get_refresh_token_url(self) -> str:
+        return self.get_oauth_access_token_url()
+
+    def get_refresh_token_params(
+        self, refresh_token: str, identity: Identity | RpcIdentity, **kwargs: Any
+    ) -> dict[str, str | None]:
+        return {"grant_type": "refresh_token", "refresh_token": refresh_token}
+
+    def get_refresh_token(
+        self, refresh_token: str, url: str, identity: Identity | RpcIdentity, **kwargs: Any
+    ) -> Response:
+        data = self.get_refresh_token_params(refresh_token, identity, **kwargs)
+
+        client_id = identity.data.get("client_id")
+        client_secret = identity.data.get("client_secret")
+        if not client_id or not client_secret:
+            raise IdentityNotValid("Missing DCR credentials")
+        headers = {"Authorization": _basic_auth_header(client_id, client_secret)}
+
+        try:
+            req = safe_urlopen(
+                url=url,
+                headers=headers,
+                data=data,
+                verify_ssl=kwargs.get("verify_ssl", True),
+            )
+            req.raise_for_status()
+        except HTTPError as e:
+            error_resp = e.response
+            exc = ApiError.from_response(error_resp, url=url)
+            if isinstance(exc, ApiUnauthorized | ApiInvalidRequestError):
+                raise IdentityNotValid from e
+            raise exc from e
+
+        return req

src/sentry/integrations/types.py

@@ -46,6 +46,7 @@ class IntegrationProviderSlug(StrEnum):
     PAGERDUTY = "pagerduty"
     OPSGENIE = "opsgenie"
     PERFORCE = "perforce"
+    DATADOG = "datadog"
     GCP = "gcp"
 
 

tests/sentry/identity/datadog/test_provider.py

@@ -7,6 +7,7 @@
 from urllib.parse import parse_qs, parse_qsl, urlparse
 
 import orjson
+import pytest
 import responses
 from django.contrib.messages.storage.fallback import FallbackStorage
 from django.contrib.sessions.backends.base import SessionBase
@@ -15,7 +16,9 @@
 from requests.exceptions import SSLError
 
 import sentry.identity
+from sentry.auth.exceptions import IdentityNotValid
 from sentry.identity.datadog.provider import (
+    DatadogIdentityProvider,
     DatadogOAuth2CallbackView,
     DatadogOAuth2DCRView,
     DatadogOAuth2LoginView,
@@ -381,3 +384,98 @@ def test_pipeline_views(self, mock_record: MagicMock) -> None:
             f"Basic {base64.b64encode(b'dcr-client-id:dcr-client-secret').decode()}"
         )
         assert exchange_token_request.headers["Authorization"] == expected_auth_header
+
+
+@control_silo_test
+class DatadogIdentityProviderTest(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self.provider = DatadogIdentityProvider()
+        self.provider.config = {"site": "datadoghq.com"}
+
+    @patch("sentry.identity.datadog.provider.get_user_info")
+    def test_build_identity(self, mock_get_user_info: MagicMock) -> None:
+        mock_get_user_info.return_value = {
+            "id": "dd-user-123",
+            "attributes": {"email": "user@example.com", "name": "Test User"},
+        }
+
+        result = self.provider.build_identity(
+            {
+                "data": {
+                    "access_token": "token-abc",
+                    "refresh_token": "refresh-xyz",
+                    "expires_in": 3600,
+                    "token_type": "Bearer",
+                    "scope": "apm_read",
+                },
+                "dcr_client_id": "dcr-client-id",
+                "dcr_client_secret": "dcr-client-secret",
+            }
+        )
+
+        assert result["id"] == "dd-user-123"
+        assert result["email"] == "user@example.com"
+        assert result["name"] == "Test User"
+        assert result["type"] == "datadog"
+        assert result["data"]["access_token"] == "token-abc"
+        assert result["data"]["refresh_token"] == "refresh-xyz"
+        assert "expires" in result["data"]
+        assert result["data"]["token_type"] == "Bearer"
+        assert result["data"]["scope"] == "apm_read"
+        assert result["data"]["client_id"] == "dcr-client-id"
+        assert result["data"]["client_secret"] == "dcr-client-secret"
+        mock_get_user_info.assert_called_once_with("token-abc", "datadoghq.com")
+
+    @patch("sentry.identity.datadog.provider.get_user_info")
+    def test_build_identity_missing_access_token(self, mock_get_user_info: MagicMock) -> None:
+        with pytest.raises(ValueError, match="did not return an access_token"):
+            self.provider.build_identity({"data": {}})
+        mock_get_user_info.assert_not_called()
+
+    @patch("sentry.identity.datadog.provider.get_user_info")
+    def test_build_identity_missing_user_attributes(self, mock_get_user_info: MagicMock) -> None:
+        mock_get_user_info.return_value = {"id": "dd-user-456", "attributes": {}}
+
+        result = self.provider.build_identity({"data": {"access_token": "token"}})
+
+        assert result["id"] == "dd-user-456"
+        assert result["email"] is None
+        assert result["name"] is None
+
+    @responses.activate
+    def test_get_refresh_token_success(self) -> None:
+        responses.add(responses.POST, TOKEN_URL, json={"access_token": "new-token"})
+
+        identity = MagicMock()
+        identity.data = {"client_id": "dcr-client", "client_secret": "dcr-secret"}
+
+        result = self.provider.get_refresh_token("refresh-token", TOKEN_URL, identity)
+
+        assert result.status_code == 200
+
+        auth_header = responses.calls[0].request.headers["Authorization"]
+        assert auth_header == f"Basic {base64.b64encode(b'dcr-client:dcr-secret').decode()}"
+
+        data = dict(parse_qsl(responses.calls[0].request.body))
+        assert data["grant_type"] == "refresh_token"
+        assert data["refresh_token"] == "refresh-token"
+        assert "client_id" not in data
+        assert "client_secret" not in data
+
+    def test_get_refresh_token_missing_dcr_credentials(self) -> None:
+        identity = MagicMock()
+        identity.data = {}
+
+        with pytest.raises(IdentityNotValid, match="Missing DCR credentials"):
+            self.provider.get_refresh_token("refresh-token", TOKEN_URL, identity)
+
+    @responses.activate
+    def test_get_refresh_token_unauthorized(self) -> None:
+        responses.add(responses.POST, TOKEN_URL, json={"error": "invalid_grant"}, status=401)
+
+        identity = MagicMock()
+        identity.data = {"client_id": "dcr-client-id", "client_secret": "dcr-client-secret"}
+
+        with pytest.raises(IdentityNotValid):
+            self.provider.get_refresh_token("refresh-token", TOKEN_URL, identity)