Update dogshell site parameter handling, fix us1 and add gov and custom url (#780)

Author: nkzou

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()

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

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