From 2f0bdcbe446a4bc49b22f7bcbc14ca5cdb0cb517 Mon Sep 17 00:00:00 2001
From: talboren <tal@keephq.dev>
Date: Wed, 14 Feb 2024 11:06:15 +0200
Subject: [PATCH] fix(datadog): oauth integration in EU (#829)

---
 keep-ui/app/providers/page.tsx                |  1 -
 keep-ui/app/providers/providers-tiles.tsx     | 10 +++++----
 .../datadog_provider/datadog_provider.py      | 21 +++++++++++++++++--
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/keep-ui/app/providers/page.tsx b/keep-ui/app/providers/page.tsx
index 9757bb301..d0852e957 100644
--- a/keep-ui/app/providers/page.tsx
+++ b/keep-ui/app/providers/page.tsx
@@ -7,7 +7,6 @@ export default async function Page({
 }: {
   searchParams?: { [key: string]: string };
 }) {
-  const session = await getServerSession(authOptions);
   return <ProvidersPage searchParams={searchParams} />;
 }
 
diff --git a/keep-ui/app/providers/providers-tiles.tsx b/keep-ui/app/providers/providers-tiles.tsx
index 4bd23bca7..74570f6ce 100644
--- a/keep-ui/app/providers/providers-tiles.tsx
+++ b/keep-ui/app/providers/providers-tiles.tsx
@@ -35,7 +35,7 @@ const ProvidersTiles = ({
   const providerName = searchParams?.get("provider_name");
 
   useEffect(() => {
-    if (providerType && providerName) {
+    if (providerType) {
       // Find the provider based on providerType and providerName
       const provider = providers.find(
         (provider) => provider.type === providerType
@@ -43,9 +43,11 @@ const ProvidersTiles = ({
 
       if (provider) {
         setSelectedProvider(provider);
-        setFormValues({
-          provider_name: providerName,
-        });
+        if (providerName) {
+          setFormValues({
+            provider_name: providerName,
+          });
+        }
         setOpenPanel(true);
       }
     }
diff --git a/keep/providers/datadog_provider/datadog_provider.py b/keep/providers/datadog_provider/datadog_provider.py
index d21bf5aa6..09ba92576 100644
--- a/keep/providers/datadog_provider/datadog_provider.py
+++ b/keep/providers/datadog_provider/datadog_provider.py
@@ -64,6 +64,15 @@ class DatadogProviderAuthConfig:
         },
         default="",
     )
+    domain: str = dataclasses.field(
+        metadata={
+            "required": True,
+            "description": "Datadog API domain",
+            "sensitive": False,
+            "hint": "https://api.datadoghq.com",
+        },
+        default="https://api.datadoghq.com",
+    )
     oauth_token: dict = dataclasses.field(
         metadata={
             "description": "For OAuth flow",
@@ -197,9 +206,14 @@ def __init__(
             self.configuration.api_key[
                 "appKeyAuth"
             ] = self.authentication_config.app_key
+            domain = self.authentication_config.domain or "https://api.datadoghq.com"
+            self.configuration.host = domain
         elif self.authentication_config.oauth_token:
+            domain = self.authentication_config.oauth_token.get(
+                "domain", "datadoghq.com"
+            )
             response = requests.post(
-                "https://api.datadoghq.com/oauth2/v1/token",
+                f"https://api.{domain}/oauth2/v1/token",
                 data={
                     "grant_type": "refresh_token",
                     "client_id": DatadogProvider.DATADOG_CLIENT_ID,
@@ -220,6 +234,7 @@ def __init__(
                 raise Exception("Could not refresh token, need to re-authenticate")
             response_json = response.json()
             self.configuration.access_token = response_json.get("access_token")
+            self.configuration.host = f"https://api.{domain}"
             # update the oauth_token refresh_token for next run
             self.config.authentication["oauth_token"]["refresh_token"] = response_json[
                 "refresh_token"
@@ -238,6 +253,7 @@ def oauth2_logic(**payload) -> dict:
         Returns:
             dict: access token to Datadog.
         """
+        domain = payload.pop("domain", "datadoghq.com")
         verifier = payload.pop("verifier", None)
         if not verifier:
             raise Exception("No verifier provided")
@@ -246,7 +262,7 @@ def oauth2_logic(**payload) -> dict:
             raise Exception("No code provided")
 
         token = requests.post(
-            "https://api.datadoghq.com/oauth2/v1/token",
+            f"https://api.{domain}/oauth2/v1/token",
             data={
                 "grant_type": "authorization_code",
                 "client_id": payload["client_id"],
@@ -267,6 +283,7 @@ def oauth2_logic(**payload) -> dict:
                 "verifier": verifier,
                 "code": code,
                 "redirect_uri": payload["redirect_uri"],
+                "domain": domain,
             }
         }