Allow optionally fetching lease data from S3

bare_metal_billing/config.py

@@ -0,0 +1,9 @@
+import os
+
+# S3 Configuration
+S3_LEASE_ENDPOINT_URL = os.getenv(
+    "S3_LEASE_ENDPOINT_URL", "https://s3.us-east-005.backblazeb2.com"
+)
+S3_LEASE_KEY_ID = os.getenv("S3_LEASE_KEY_ID")
+S3_LEASE_APP_KEY = os.getenv("S3_LEASE_APP_KEY")
+S3_LEASE_BUCKET = os.getenv("S3_LEASE_BUCKET")

bare_metal_billing/main.py

@@ -6,7 +6,7 @@
 
 from nerc_rates import load_from_url
 
-from bare_metal_billing import models, billing
+from bare_metal_billing import models, billing, s3_bucket, config
 
 
 logging.basicConfig(level=logging.INFO)
@@ -69,7 +69,9 @@ def main():
     )
     parser.add_argument(
         "bm_usage_file",
-        help="Bare Metal usage JSON file to be processed",
+        nargs="?",
+        default=None,
+        help="Bare Metal usage JSON file to be processed. If not provided, defaults to fetching from S3.",
     )
     parser.add_argument(
         "--rate-fc430-su", default=0, type=Decimal, help="Rate of FC430 SU/hr"
@@ -107,6 +109,16 @@ def main():
     logger.info(f"Interval for processing {args.start} - {args.end}.")
     logger.info(f"Invoice file will be saved to {args.output_file}.")
 
+    # Default to fetching lease file one day before billing end date because end date is not-inclusive
+    if args.bm_usage_file:
+        bm_usage_file = args.bm_usage_file
+    else:
+        previous_day = (default_end_argument() - timedelta(days=1)).strftime("%Y%m%d")
+        bm_usage_file = s3_bucket.fetch_s3(
+            config.S3_LEASE_BUCKET,
+            f"{args.invoice_month}/esi-lease-{previous_day}.json",
+        )
+
     su_rates_dict = {}
     if args.use_nerc_rates:
         nerc_repo_rates = load_from_url()
@@ -125,7 +137,7 @@ def main():
         su_rates_dict,
     )
 
-    with open(args.bm_usage_file, "r") as f:
+    with open(bm_usage_file, "r") as f:
         input_bm_json = json.load(f)
 
     input_invoice = models.BMUsageData.model_validate(input_bm_json)

bare_metal_billing/s3_bucket.py

@@ -0,0 +1,34 @@
+import os
+import logging
+import functools
+
+import boto3
+
+from bare_metal_billing import config
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+@functools.lru_cache
+def get_bucket(bucket_name: str):
+    if not config.S3_LEASE_APP_KEY or not config.S3_LEASE_KEY_ID:
+        raise RuntimeError(
+            "Please set the environment variables S3_LEASE_APP_KEY, S3_LEASE_KEY_ID"
+        )
+
+    s3_resource = boto3.resource(
+        service_name="s3",
+        endpoint_url=config.S3_LEASE_ENDPOINT_URL,
+        aws_access_key_id=config.S3_LEASE_KEY_ID,
+        aws_secret_access_key=config.S3_LEASE_APP_KEY,
+    )
+
+    return s3_resource.Bucket(bucket_name)
+
+
+def fetch_s3(bucket_name: str, s3_filepath: str) -> str:
+    local_name = os.path.basename(s3_filepath)
+    lease_bucket = get_bucket(bucket_name)
+    lease_bucket.download_file(s3_filepath, local_name)
+    return local_name

bare_metal_billing/tests/test_s3_bucket.py

@@ -0,0 +1,21 @@
+from unittest import TestCase, mock
+
+from bare_metal_billing import s3_bucket
+
+
+class TestS3Bucket(TestCase):
+    def test_fetch_s3(self):
+        mock_bucket = mock.MagicMock()
+        mock_get_bucket = mock.MagicMock(return_value=mock_bucket)
+
+        with mock.patch("bare_metal_billing.s3_bucket.get_bucket", mock_get_bucket):
+            test_s3_filepath = "path/to/testfile.json"
+            expected_local_name = "testfile.json"
+
+            local_name = s3_bucket.fetch_s3("foo-bucket", test_s3_filepath)
+
+            self.assertEqual(local_name, expected_local_name)
+            mock_get_bucket.assert_called_once()
+            mock_bucket.download_file.assert_called_once_with(
+                test_s3_filepath, expected_local_name
+            )

requirements.txt

@@ -1,2 +1,3 @@
 git+https://github.com/CCI-MOC/nerc-rates@d3bcfb2#egg=nerc_rates
 pydantic
+boto3