From 86184cc20b2b509598fe34fccc8cc7b5e7c02a24 Mon Sep 17 00:00:00 2001
From: devsjc <47188100+devsjc@users.noreply.github.com>
Date: Tue, 17 Oct 2023 15:39:44 +0100
Subject: [PATCH] Enable pulling 84 hours of India data (#24)

---
 nwp/jobs.py    | 24 +++++++++++++++++-------
 pyproject.toml |  2 +-
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/nwp/jobs.py b/nwp/jobs.py
index c585889..3e52189 100644
--- a/nwp/jobs.py
+++ b/nwp/jobs.py
@@ -84,19 +84,29 @@ class NWPConsumerDagDefinition:
     """A class to define the NWPConsumerDagDefinition."""
 
     def __init__(
-        self, area: str, source: str, storage_path: str | None = None
+        self, source: str, storage_path: str | None = None, env_overrides: dict[str, str] | None = None
     ) -> "NWPConsumerDagDefinition":
         """Create a NWPConsumerDagDefinition."""
-        self.area = area
         self.source = source
+        area = env_overrides.get("ECMWF_AREA", "no-area")
         self.storage_path = \
             storage_path or \
-            f'/mnt/storage_b/data/ocf/solar_pv_nowcasting/nowcasting_dataset_pipeline/NWP/ECMWF/{self.area}'
+            f'/mnt/storage_b/data/ocf/solar_pv_nowcasting/nowcasting_dataset_pipeline/NWP/ECMWF/{area}'
+        self.env_overrides = env_overrides
 
 nwp_consumer_jobs: dict[str, NWPConsumerDagDefinition] = {
-    "uk": NWPConsumerDagDefinition(area="uk", source="ecmwf-mars"),
-    "india": NWPConsumerDagDefinition(area="nw-india", source="ecmwf-mars"),
-    "malta": NWPConsumerDagDefinition(area="malta", source="ecmwf-mars")
+    "uk": NWPConsumerDagDefinition(
+        source="ecmwf-mars",
+        env_overrides={"ECMWF_AREA": "uk"}
+    ),
+    "india": NWPConsumerDagDefinition(
+        source="ecmwf-mars",
+        env_overrides={"ECMWF_AREA": "nw-india", "ECMWF_HOURS": "84"}
+    ),
+    "malta": NWPConsumerDagDefinition(
+        source="ecmwf-mars",
+        env_overrides={"ECMWF_AREA": "malta"}
+    ),
 }
 
 
@@ -110,7 +120,7 @@ def partitioned_config_func(partition_key: str) -> dict[str, Any]:
             date_from=time_window.start.strftime("%Y-%m-%d"),
             date_to=time_window.start.strftime("%Y-%m-%d"),
             source=dagdef.source,
-            env_overrides={"ECMWF_AREA": dagdef.area},
+            env_overrides=dagdef.env_overrides,
             zarr_dir=f"{dagdef.storage_path}/zarr",
             raw_dir=f"{dagdef.storage_path}/raw",
         )
diff --git a/pyproject.toml b/pyproject.toml
index 2dbdfa9..009426b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,7 @@ dependencies = [
     "dagster-docker == 0.20.11",
     "huggingface-hub == 0.16.4",
     "numpy >= 1.23.0",
-    "nwp-consumer == 0.1.21",
+    "nwp-consumer == 0.1.23",
     "ocf-blosc2 == 0.0.3",
     "pathlib == 1.0.1",
     "requests >= 2.28.0",