Merge branch 'add-api-data-host-option' into 'master'

Add option `data_host` to /ws/data endpoint.

See merge request !278

geomagio/api/ws/DataApiQuery.py

@@ -40,6 +40,18 @@ class SamplingPeriod(float, enum.Enum):
     DAY = 86400.0
 
 
+class DataHost(str, enum.Enum):
+    # recognized public Edge data hosts, plus one user-specified
+    DEFAULT = os.getenv("DATA_HOST", "edgecwb.usgs.gov")
+    EDGECWB = "edgecwb.usgs.gov"
+    CWBPUB = "cwbpub.cr.usgs.gov"
+    CWBPUB2 = "cwbp2.cr.usgs.gov"
+
+    @classmethod
+    def values(cls) -> List[str]:
+        return [t.value for t in cls]
+
+
 class DataApiQuery(BaseModel):
     id: str
     starttime: UTCDateTime = None
@@ -47,7 +59,8 @@ class DataApiQuery(BaseModel):
     elements: List[str] = DEFAULT_ELEMENTS
     sampling_period: SamplingPeriod = SamplingPeriod.MINUTE
     data_type: Union[DataType, str] = DataType.VARIATION
-    format: OutputFormat = OutputFormat.IAGA2002
+    format: Union[OutputFormat, str] = OutputFormat.IAGA2002
+    data_host: Union[DataHost, str] = DataHost.DEFAULT
 
     @validator("data_type")
     def validate_data_type(
@@ -60,6 +73,17 @@ class DataApiQuery(BaseModel):
             )
         return data_type
 
+    @validator("data_host")
+    def validate_data_host(
+        cls, data_host: Union[DataHost, str]
+    ) -> Union[DataHost, str]:
+        if data_host not in DataHost.values():
+            raise ValueError(
+                # don't advertise acceptable hosts
+                f"Bad data_host value '{data_host}'."
+            )
+        return data_host
+
     @validator("elements", pre=True, always=True)
     def validate_elements(cls, elements: List[str]) -> List[str]:
         if not elements:

geomagio/api/ws/app.py

@@ -28,6 +28,9 @@ app = FastAPI(
     docs_url="/docs",
     root_path="/ws",
     title="Geomagnetism Data Web Service",
+    swagger_ui_parameters={
+        "defaultModelsExpandDepth": 0,
+    },
 )
 
 app.include_router(algorithms.router)

geomagio/api/ws/data.py

@@ -16,6 +16,7 @@ from .DataApiQuery import (
     DataType,
     OutputFormat,
     SamplingPeriod,
+    DataHost,
 )
 
 
@@ -29,7 +30,7 @@ def get_data_factory(
     data_factory
         Edge or miniseed factory object
     """
-    host = os.getenv("DATA_HOST", "edgecwb.usgs.gov")
+    host = query.data_host or DataHost.DEFAULT
     sampling_period = query.sampling_period
     if query.id in ASL_OBSERVATORY_INDEX:
         factory = IRISFactory(network="IU", locationCode="40")
@@ -81,7 +82,16 @@ def get_data_query(
         " NOTE: the USGS web service also supports specific EDGE location codes."
         " For example: R0 is 'internet variation'",
     ),
-    format: OutputFormat = Query(OutputFormat.IAGA2002),
+    format: Union[OutputFormat, str] = Query(
+        OutputFormat.IAGA2002,
+        title="Output Format",
+        description="Format of returned time series data.",
+    ),
+    data_host: Union[DataHost, str] = Query(
+        DataHost.DEFAULT,
+        title="Data Host",
+        description="Edge host to pull data from.",
+    ),
 ) -> DataApiQuery:
     """Define query parameters used for webservice requests.
 
@@ -114,6 +124,7 @@ def get_data_query(
         "sampling_period",
         "type",
         "format",
+        "data_host",
     ]
     invalid_params = []
     for param in request.query_params.keys():
@@ -131,6 +142,7 @@ def get_data_query(
         sampling_period=sampling_period,
         data_type=data_type,
         format=format,
+        data_host=data_host,
     )
     return query