feat: allow to override parquet_prefix and include idle costs

src/opencost_parquet_exporter.py

@@ -37,6 +37,7 @@ def get_config(
         window_end=None,
         s3_bucket=None,
         file_key_prefix=None,
+        parquet_prefix=None,
         aggregate_by=None,
         step=None,
         resolution=None,
@@ -152,10 +153,19 @@ def get_config(
             datetime.now() - timedelta(1), '%Y-%m-%d')
         window_start = yesterday+'T00:00:00Z'
         window_end = yesterday+'T23:59:59Z'
+    if parquet_prefix is None:
+      window = datetime.strptime(window_start, "%Y-%m-%dT%H:%M:%SZ")
+      parquet_prefix = os.environ.get(
+        'OPENCOST_PARQUET_PREFIX',
+        f"{file_key_prefix}/year={window.year}/month={window.month}/day={window.day}"
+      )
+
+    config['parquet_prefix'] = parquet_prefix
     window = f"{window_start},{window_end}"
     config['window_start'] = window_start
     config['params'] = [
         ("window", window),
+        ("aggregate", aggregate_by),
         ("includeIdle", include_idle),
         ("idleByNode", idle_by_node),
         ("includeProportionalAssetResourceCosts", "false"),
@@ -279,7 +289,6 @@ def save_result(processed_result, config):
 
 # pylint: disable=C0116
 
-
 def main():
     # TODO: Error handling when load fails
     print("Starting run")

src/test_opencost_parquet_exporter.py

@@ -3,6 +3,7 @@
 from unittest.mock import patch, MagicMock, mock_open
 import json
 import os
+from datetime import datetime
 import requests
 from freezegun import freeze_time
 from opencost_parquet_exporter import get_config, request_data, load_config_file
@@ -139,6 +140,38 @@ def test_get_config_no_window_end(self):
             self.assertNotIn('s3_bucket', config)
             self.assertEqual(config['params'][0][1], window)
 
+    def test_get_config_with_idle_env_vars(self):
+      """Test get_config overrides parquet prefix when env var is set."""
+      with patch.dict(os.environ, {
+        'OPENCOST_PARQUET_INCLUDE_IDLE': 'true',
+        'OPENCOST_PARQUET_IDLE_BY_NODE': 'true'}, clear=True):
+        config = get_config()
+        self.assertEqual(config['params'][2][0], 'includeIdle')
+        self.assertEqual(config['params'][2][1], 'true')
+
+        self.assertEqual(config['params'][3][0], 'idleByNode')
+        self.assertEqual(config['params'][3][1], 'true')
+
+    @freeze_time('2024-06-29')
+    def test_get_config_with_parquet_prefix_env_var(self):
+
+      yesterday = '2024-06-28'
+      window_start = yesterday+'T01:00:00Z'
+      window = datetime.strptime(window_start, "%Y-%m-%dT%H:%M:%SZ")
+
+      """Test get_config overrides parquet prefix when env var is set."""
+      with patch.dict(os.environ, {
+        'OPENCOST_PARQUET_FILE_KEY_PREFIX': 'test-prefix/',
+        'OPENCOST_PARQUET_PREFIX': f"year={window.year}/month={window.month}/day={window.day}/hour={window.hour}",
+        'OPENCOST_PARQUET_STEP': '1m'}, clear=True):
+        config = get_config()
+        self.assertEqual(config['parquet_prefix'], 'year=2024/month=6/day=28/hour=1')
+
+    @freeze_time('2024-06-29')
+    def test_get_config_with_default_parquet_prefix_env_var(self):
+      """Test get_config with default parquet prefix when env var is not set."""
+      config = get_config()
+      self.assertEqual(config['parquet_prefix'], '/tmp//year=2024/month=6/day=28')
 
 class TestRequestData(unittest.TestCase):
     """ Test request_data method """