Merge branch 'master' into kevinzou/sanitize_dogshell_inputs

DataDog · Jul 11, 2023 · 2a49f88 · 2a49f88
2 parents 57f3a78 + 8ae1cf6
commit 2a49f88
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 36 deletions.
diff --git a/datadog/dogshell/__init__.py b/datadog/dogshell/__init__.py
@@ -74,20 +74,8 @@ def main():
     parser.add_argument(
         "--api_host",
         help="Datadog site to send data, us (datadoghq.com), eu (datadoghq.eu), us3 (us3.datadoghq.com), \
-              us5 (us5.datadoghq.com), or ap1 (ap1.datadoghq.com). default: us",
+              us5 (us5.datadoghq.com), ap1 (ap1.datadoghq.com), gov (ddog-gov.com), or custom url. default: us",
         dest="api_host",
-        choices=[
-            "datadoghq.com",
-            "us",
-            "datadoghq.eu",
-            "eu",
-            "us3.datadoghq.com",
-            "us3",
-            "us5.datadoghq.com",
-            "us5",
-            "ap1.datadoghq.com",
-            "ap1"
-        ],
     )
 
     config = DogshellConfig()

diff --git a/datadog/dogshell/common.py b/datadog/dogshell/common.py
@@ -47,17 +47,20 @@ def load(self, config_file, api_key, app_key, api_host):
         config = configparser.ConfigParser()
 
         if api_host is not None:
-            if api_host in ("us" or "datadoghq.com"):
-                self["api_host"] = "https://datadoghq.com"
+            if api_host in ("datadoghq.com", "us"):
+                self["api_host"] = "https://api.datadoghq.com"
             elif api_host in ("datadoghq.eu", "eu"):
-                self["api_host"] = "https://datadoghq.eu"
+                self["api_host"] = "https://api.datadoghq.eu"
             elif api_host in ("us3.datadoghq.com", "us3"):
-                self["api_host"] = "https://us3.datadoghq.com"
+                self["api_host"] = "https://api.us3.datadoghq.com"
             elif api_host in ("us5.datadoghq.com", "us5"):
-                self["api_host"] = "https://us5.datadoghq.com"
+                self["api_host"] = "https://api.us5.datadoghq.com"
             elif api_host in ("ap1.datadoghq.com", "ap1"):
-                self["api_host"] = "https://ap1.datadoghq.com"
-
+                self["api_host"] = "https://api.ap1.datadoghq.com"
+            elif api_host in ("ddog-gov.com", "gov"):
+                self["api_host"] = "https://api.ddog-gov.com"
+            else:
+                self["api_host"] = api_host
         if api_key is not None and app_key is not None:
             self["api_key"] = api_key
             self["app_key"] = app_key

diff --git a/datadog/dogshell/wrap.py b/datadog/dogshell/wrap.py
@@ -284,22 +284,11 @@ def parse_options(raw_args=None):
         "-s",
         "--site",
         action="store",
-        type="choice",
+        type="string",
         default="datadoghq.com",
-        choices=[
-            "datadoghq.com",
-            "us",
-            "datadoghq.eu",
-            "eu",
-            "us3.datadoghq.com",
-            "us3",
-            "us5.datadoghq.com",
-            "us5",
-            "ap1.datadoghq.com",
-            "ap1"
-        ],
         help="The site to send data. Accepts us (datadoghq.com), eu (datadoghq.eu), \
-us3 (us3.datadoghq.com), us5 (us5.datadoghq.com), or ap1 (ap1.datadoghq.com). default: us",
+us3 (us3.datadoghq.com), us5 (us5.datadoghq.com), or ap1 (ap1.datadoghq.com), \
+gov (ddog-gov.com), or custom url. default: us",
     )
     parser.add_option(
         "-m",
@@ -430,16 +419,20 @@ def main():
         options.buffer_outs,
     )
 
-    if options.site in ("datadoghq.eu", "eu"):
+    if options.site in ("datadoghq.com", "us"):
+        api_host = "https://api.datadoghq.com"
+    elif options.site in ("datadoghq.eu", "eu"):
         api_host = "https://api.datadoghq.eu"
     elif options.site in ("us3.datadoghq.com", "us3"):
         api_host = "https://api.us3.datadoghq.com"
     elif options.site in ("us5.datadoghq.com", "us5"):
         api_host = "https://api.us5.datadoghq.com"
     elif options.site in ("ap1.datadoghq.com", "ap1"):
         api_host = "https://api.ap1.datadoghq.com"
+    elif options.site in ("ddog-gov.com", "gov"):
+        api_host = "https://api.ddog-gov.com"
     else:
-        api_host = "https://api.datadoghq.com"
+        api_host = options.site
 
     initialize(api_key=options.api_key, api_host=api_host)
     host = api._host_name