diff --git a/windwatts-api/app/config/sample_windwatts_data_config.json b/windwatts-api/app/config/sample_windwatts_data_config.json index 1d893c7..e9e17e7 100644 --- a/windwatts-api/app/config/sample_windwatts_data_config.json +++ b/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" -} \ No newline at end of file + "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"] } + } + } +} diff --git a/windwatts-api/app/config_manager.py b/windwatts-api/app/config_manager.py index c2a8f17..8e652a1 100644 --- a/windwatts-api/app/config_manager.py +++ b/windwatts-api/app/config_manager.py @@ -64,7 +64,7 @@ def _get_config_from_env(self): # Scan for all SOURCES__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(): @@ -72,11 +72,12 @@ def _get_config_from_env(self): 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: diff --git a/windwatts-api/app/schemas.py b/windwatts-api/app/schemas.py index 6f845ea..a7c22a3 100644 --- a/windwatts-api/app/schemas.py +++ b/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" ) - diff --git a/windwatts-api/app/spatial/global_spatial_manager.py b/windwatts-api/app/spatial/global_spatial_manager.py index 6d226e4..c29ec83 100644 --- a/windwatts-api/app/spatial/global_spatial_manager.py +++ b/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 diff --git a/windwatts-api/app/utils/athena_query_client.py b/windwatts-api/app/utils/athena_query_client.py index 63b4e16..c2ab362 100644 --- a/windwatts-api/app/utils/athena_query_client.py +++ b/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"))) \ No newline at end of file + return pd.read_csv(StringIO(obj["Body"].read().decode("utf-8")))