diff --git a/bare_metal_billing/config.py b/bare_metal_billing/config.py new file mode 100644 index 0000000..c7fafd7 --- /dev/null +++ b/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") diff --git a/bare_metal_billing/main.py b/bare_metal_billing/main.py index 2b82025..05c27af 100644 --- a/bare_metal_billing/main.py +++ b/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) diff --git a/bare_metal_billing/s3_bucket.py b/bare_metal_billing/s3_bucket.py new file mode 100644 index 0000000..41d30af --- /dev/null +++ b/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 diff --git a/bare_metal_billing/tests/test_s3_bucket.py b/bare_metal_billing/tests/test_s3_bucket.py new file mode 100644 index 0000000..2a3a3e6 --- /dev/null +++ b/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 + ) diff --git a/requirements.txt b/requirements.txt index 0db7587..fc5095f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ git+https://github.com/CCI-MOC/nerc-rates@d3bcfb2#egg=nerc_rates pydantic +boto3