From 2e6aa125d9f712088f49b2b304023995be75322e Mon Sep 17 00:00:00 2001
From: Slesa Adhikari <slesaad@gmail.com>
Date: Mon, 9 Oct 2023 16:04:04 -0500
Subject: [PATCH] Header host override for stac and raster apis (#47)

---
 .example.env                           |  4 +--
 raster_api/infrastructure/config.py    | 30 ++++++++++++++++++++---
 raster_api/infrastructure/construct.py | 11 ++++++---
 stac_api/infrastructure/config.py      | 34 ++++++++++++++++++++------
 4 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/.example.env b/.example.env
index 2305c8d4..1a14613f 100644
--- a/.example.env
+++ b/.example.env
@@ -21,9 +21,7 @@ VEDA_DOMAIN_ALT_HOSTED_ZONE_NAME=[OPTIONAL SECOND DOMAIN]
 VEDA_RASTER_ENABLE_MOSAIC_SEARCH=TRUE
 VEDA_RASTER_DATA_ACCESS_ROLE_ARN=[OPTIONAL ARN OF IAM ROLE TO BE ASSUMED BY RASTER API] 
 
-
 VEDA_RASTER_PATH_PREFIX=[OPTIONAL PATH PREFIX TO ADD TO TITILER ENDPOINTS]
 VEDA_STAC_PATH_PREFIX=[OPTIONAL PATH PREFIX TO ADD TO TITILER ENDPOINTS]
-VEDA_STAC_HOST=[OPTIONAL HOST FOR THE STAC API]
 
-VEDA_DB_RDS_TYPE=[OPTIONAL RDS DATABASE TYPE, ex. t3.small, r5.16xlarge, etc]
+VEDA_HOST=[OPTIONAL HOST/DOMAIN_NAME TO PROPAGATE TO STAC AND RASTER APIS]
diff --git a/raster_api/infrastructure/config.py b/raster_api/infrastructure/config.py
index 6562848b..6d41cfa5 100644
--- a/raster_api/infrastructure/config.py
+++ b/raster_api/infrastructure/config.py
@@ -6,8 +6,19 @@
 from pydantic import BaseSettings, Field
 
 
+class MyConfig(BaseSettings.Config):
+    """Custom config class that support multiple env_prefixes"""
+
+    @classmethod
+    def prepare_field(cls, field) -> None:
+        """Workaround to not overwrite ENV_PREFIX"""
+        if "env_names" in field.field_info.extra:
+            return
+        return super().prepare_field(field)
+
+
 class vedaRasterSettings(BaseSettings):
-    """Application settings"""
+    """Raster settings"""
 
     # Default options are optimized for CloudOptimized GeoTIFF
     # For more information on GDAL env see: https://gdal.org/user/configoptions.html
@@ -68,11 +79,24 @@ class vedaRasterSettings(BaseSettings):
         description="Optional path prefix to add to all api endpoints",
     )
 
-    class Config:
+    class Config(MyConfig):
         """model config"""
 
         env_file = ".env"
         env_prefix = "VEDA_RASTER_"
 
 
-veda_raster_settings = vedaRasterSettings()
+class Settings(vedaRasterSettings):
+    """Application Settings"""
+
+    host: Optional[str] = Field(
+        "",
+        description="Optional host to send to raster api",  # propagate cf url to raster api
+    )
+
+    class Config(MyConfig):
+        "Model config"
+        env_prefix = "VEDA_"
+
+
+veda_raster_settings = Settings()
diff --git a/raster_api/infrastructure/construct.py b/raster_api/infrastructure/construct.py
index a0795f0f..e9dbba22 100644
--- a/raster_api/infrastructure/construct.py
+++ b/raster_api/infrastructure/construct.py
@@ -79,10 +79,15 @@ def __init__(
                 "AWS_REQUEST_PAYER", veda_raster_settings.aws_request_payer
             )
 
-        raster_api_integration = (
-            aws_apigatewayv2_integrations_alpha.HttpLambdaIntegration(
-                construct_id, veda_raster_function
+        raster_api_integration = aws_apigatewayv2_integrations_alpha.HttpLambdaIntegration(
+            construct_id,
+            handler=veda_raster_function,
+            parameter_mapping=aws_apigatewayv2_alpha.ParameterMapping().overwrite_header(
+                "host",
+                aws_apigatewayv2_alpha.MappingValue.custom(veda_raster_settings.host),
             )
+            if veda_raster_settings.host
+            else None,
         )
 
         domain_mapping = None
diff --git a/stac_api/infrastructure/config.py b/stac_api/infrastructure/config.py
index 7b8db3e6..e38ee1d3 100644
--- a/stac_api/infrastructure/config.py
+++ b/stac_api/infrastructure/config.py
@@ -4,8 +4,19 @@
 from pydantic import BaseSettings, Field
 
 
+class MyConfig(BaseSettings.Config):
+    """Custom config class that support multiple env_prefixes"""
+
+    @classmethod
+    def prepare_field(cls, field) -> None:
+        """Workaround to not overwrite ENV_PREFIX"""
+        if "env_names" in field.field_info.extra:
+            return
+        return super().prepare_field(field)
+
+
 class vedaSTACSettings(BaseSettings):
-    """Application settings"""
+    """STAC settings"""
 
     env: Dict = {}
 
@@ -23,16 +34,25 @@ class vedaSTACSettings(BaseSettings):
         description="Optional path prefix to add to all api endpoints",
     )
 
+    class Config(MyConfig):
+        """model config"""
+
+        env_file = ".env"
+        env_prefix = "VEDA_STAC_"
+
+
+class Settings(vedaSTACSettings):
+    """Application Settings"""
+
     host: Optional[str] = Field(
         "",
         description="Optional host to send to stac api",  # stac api populates the urls in the catalog based on this
     )
 
-    class Config:
-        """model config"""
-
-        env_file = ".env"
-        env_prefix = "VEDA_STAC_"
+    class Config(MyConfig):
+        "Model config"
+        env_prefix = "VEDA_"
 
 
-veda_stac_settings = vedaSTACSettings()
+veda_stac_settings = Settings()
+print(veda_stac_settings)