Implement automatic pagination for EIA data fetching

Author: mahendra-918

src/open_data_pvnet/scripts/fetch_eia_data.py

@@ -77,25 +77,45 @@ def get_data(
                 else:
                     params[f"facets[{key}][]"] = value
 
+        all_data = []
+        
         try:
-            logger.info(f"Fetching data from {url}...")
-            response = requests.get(url, params=params)
-            response.raise_for_status()
+            current_offset = offset
+            while True:
+                # Create a fresh copy of params for each request to avoid mutating history
+                request_params = params.copy()
+                request_params["offset"] = current_offset
+                
+                logger.info(f"Fetching data from {url}, offset={current_offset}...")
+                response = requests.get(url, params=request_params)
+                response.raise_for_status()
+                
+                payload = response.json()
+                if "response" in payload and "data" in payload["response"]:
+                    data = payload["response"]["data"]
+                    if not data:
+                        logger.info("No more data returned from API.")
+                        break
+                    
+                    all_data.extend(data)
+                    
+                    if len(data) < length:
+                        break
+                        
+                    current_offset += length
+                else:
+                    logger.error(f"Unexpected API response format: {payload.keys()}")
+                    break
 
-            payload = response.json()
-            if "response" in payload and "data" in payload["response"]:
-                data = payload["response"]["data"]
-                if not data:
-                    logger.warning("No data returned from API.")
-                    return None
-                return pd.DataFrame(data)
-            else:
-                logger.error(f"Unexpected API response format: {payload.keys()}")
+            if not all_data:
+                logger.warning("No data retrieved.")
                 return None
+                
+            return pd.DataFrame(all_data)
 
         except requests.exceptions.RequestException as e:
             logger.error(f"Request failed: {e}")
-            if response is not None:
+            if 'response' in locals() and response is not None:
                 logger.error(f"Response: {response.text}")
             return None
 

tests/test_eia.py

@@ -107,3 +107,40 @@ def test_get_dataset_success(mock_response):
         assert "datetime_gmt" in ds.coords or "datetime_gmt" in ds.indexes
         assert "value" in ds.data_vars
         assert len(ds.datetime_gmt) == 2
+
+def test_get_data_pagination():
+    page1 = {
+        "response": {
+            "data": [
+                {"period": "2023-01-01T00", "value": 100},
+                {"period": "2023-01-01T01", "value": 150},
+            ]
+        }
+    }
+    page2 = {
+        "response": {
+            "data": [
+                {"period": "2023-01-01T02", "value": 200},
+            ]
+        }
+    }
+    
+    mock_resp1 = Mock()
+    mock_resp1.json.return_value = page1
+    mock_resp1.raise_for_status.return_value = None
+    
+    mock_resp2 = Mock()
+    mock_resp2.json.return_value = page2
+    mock_resp2.raise_for_status.return_value = None
+    
+    with patch("requests.get", side_effect=[mock_resp1, mock_resp2]) as mock_get:
+        eia = EIAData(api_key="test_key")
+        
+        df = eia.get_data("route", "start", "end", length=2)
+        
+        assert len(df) == 3
+        assert mock_get.call_count == 2
+        
+        call_args_list = mock_get.call_args_list
+        assert call_args_list[0][1]["params"]["offset"] == 0
+        assert call_args_list[1][1]["params"]["offset"] == 2