fix: AWS Source Key

windwatts-api/app/config/sample_windwatts_data_config.json

@@ -1,10 +1,27 @@
 {
   "region_name": "us-east-1",
-  "bucket_name": "sample-bucket",
-  "database": "sample_database",
   "output_location": "s3://sample-output-location/",
   "output_bucket": "sample-output-bucket",
-  "athena_table_name": "sample_table",
-  "alt_athena_table_name": "sample_alt_table",
-  "athena_workgroup": "sample_workgroup"
-}
+  "database": "sample_database",
+  "athena_workgroup": "sample_workgroup",
+  "sources": {
+    "wtk-timeseries": {
+      "bucket_name": "sample-bucket",
+      "athena_table_name": "sample_table",
+      "alt_athena_table_name": "sample_alt_table",
+      "capabilities": { "avg_types": ["all", "annual", "monthly", "hourly"] }
+    },
+    "era5-quantiles": {
+      "bucket_name": "sample-bucket",
+      "athena_table_name": "sample_table",
+      "alt_athena_table_name": "sample_alt_table",
+      "capabilities": { "avg_types": ["all", "annual"] }
+    },
+    "ensemble-quantiles": {
+      "bucket_name": "sample-bucket",
+      "athena_table_name": "sample_table",
+      "alt_athena_table_name": "sample_alt_table",
+      "capabilities": { "avg_types": ["all"] }
+    }
+  }
+}

windwatts-api/app/config_manager.py

@@ -64,19 +64,20 @@ def _get_config_from_env(self):
         # Scan for all SOURCES_<SOURCE>_FIELD_NAME env vars
         sources = {}
         prefix = "SOURCES_"
-        suffixes = ["_BUCKET_NAME", "_ATHENA_TABLE_NAME", "_ALT_ATHENA_TABLE_NAME"]
+        suffixes = ["_ALT_ATHENA_TABLE_NAME", "_ATHENA_TABLE_NAME", "_BUCKET_NAME"]
         env = os.environ
         source_fields = {}
         for key, value in env.items():
             if key.startswith(prefix):
                 rest = key[len(prefix) :]
                 for suffix in suffixes:
                     if rest.endswith(suffix):
-                        source = rest[: -len(suffix)].lower()
+                        source = rest[: -len(suffix)].lower().replace("_", "-")
                         field = suffix[1:].lower()  # e.g. 'bucket_name'
                         if source not in source_fields:
                             source_fields[source] = {}
                         source_fields[source][field] = value
+                        break
         # Package the sources with required fields into `sources`
         for source, fields in source_fields.items():
             if "bucket_name" in fields and "athena_table_name" in fields:

windwatts-api/app/schemas.py

@@ -710,10 +710,15 @@ class ProductionRequestPayload(BaseModel):
         }
     }
 
+
 class AthenaSourceConfig(BaseModel):
     bucket_name: str = Field(..., description="S3 bucket where the dataset lives")
-    athena_table_name: str = Field(..., description="Primary Athena table (partitioned by index)")
-    alt_athena_table_name: str = Field("", description="Alternative table for non-index queries")
+    athena_table_name: str = Field(
+        ..., description="Primary Athena table (partitioned by index)"
+    )
+    alt_athena_table_name: str = Field(
+        "", description="Alternative table for non-index queries"
+    )
     capabilities: Optional[Dict[str, List[str]]] = Field(
         None, description="Optional capabilities like supported avg_types"
     )
@@ -728,4 +733,3 @@ class AthenaConfig(BaseModel):
     sources: Dict[str, AthenaSourceConfig] = Field(
         ..., description="Map of model_key to source config"
     )
-

windwatts-api/app/spatial/global_spatial_manager.py

@@ -15,6 +15,7 @@
 
 _initialized = False
 
+
 def init_spatial():
     "Load grids and register lookups for all models in MODEL_CONFIG"
     global _initialized

windwatts-api/app/utils/athena_query_client.py

@@ -22,7 +22,9 @@ def __init__(self, config: AthenaConfig, source: AthenaSourceConfig):
             read_timeout=5,
             retries={"max_attempts": 2, "mode": "standard"},
         )
-        self._athena = boto3.client("athena", region_name=config.region_name, config=boto_cfg)
+        self._athena = boto3.client(
+            "athena", region_name=config.region_name, config=boto_cfg
+        )
         self._s3 = boto3.client("s3", region_name=config.region_name, config=boto_cfg)
 
     def query(self, grid_idx: str) -> pd.DataFrame:
@@ -66,4 +68,4 @@ def _execute(self, query: str) -> pd.DataFrame:
         output = resp["QueryExecution"]["ResultConfiguration"]["OutputLocation"]
         bucket, key = output.replace("s3://", "").split("/", 1)
         obj = self._s3.get_object(Bucket=bucket, Key=key)
-        return pd.read_csv(StringIO(obj["Body"].read().decode("utf-8")))
+        return pd.read_csv(StringIO(obj["Body"].read().decode("utf-8")))