From 15189860607e7bd9599c5f678acb543fbcdda9dc Mon Sep 17 00:00:00 2001
From: Quan Pham <qmpham2019@gmail.com>
Date: Thu, 13 Jun 2024 11:41:59 +0700
Subject: [PATCH] Implemented 'nerc-rates' for storage invoicing

---
 requirements.txt                              |  1 +
 setup.cfg                                     |  1 +
 .../commands/calculate_storage_gb_hours.py    | 19 +++++++++++++++++--
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 2a3ce5af..bd9e5904 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+git+https://github.com/CCI-MOC/nerc-rates@74eb4a7#egg=nerc_rates
 boto3
 coldfront >= 1.0.4
 python-cinderclient  # TODO: Set version for OpenStack Clients
diff --git a/setup.cfg b/setup.cfg
index 7231f875..f24c34f8 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -20,6 +20,7 @@ package_dir =
 packages = find:
 python_requires = >=3.8
 install_requires =
+    nerc_rates @ git+https://github.com/CCI-MOC/nerc-rates@74eb4a7
     boto3
     coldfront >= 1.0.4
     python-cinderclient
diff --git a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py
index 33d07264..f61f6624 100644
--- a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py
+++ b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py
@@ -14,6 +14,8 @@
 from coldfront.core.allocation.models import Allocation, AllocationStatusChoice
 import pytz
 
+from nerc_rates import load_from_url
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -187,6 +189,19 @@ def process_invoice_row(allocation, attrs, su_name, rate):
             resources__in=openshift_resources
         )
 
+        rates = load_from_url()
+        openstack_storage_rate = openshift_storage_rate = Decimal(
+            rates.get_value_at('Storage GB Rate', options["invoice_month"]))
+        
+        if options['openstack_gb_rate']:
+            openstack_storage_rate = options['openstack_gb_rate']
+
+        if options['openshift_gb_rate']:
+            openshift_storage_rate = options['openshift_gb_rate']
+
+        logger.info(f'Using storage rate {openstack_storage_rate} (Openstack) and 
+                    {openshift_storage_rate} (Openshift) for {options["invoice_month"]}')
+
         logger.info(f'Writing to {options["output"]}.')
         with open(options['output'], 'w', newline='') as f:
             csv_invoice_writer = csv.writer(
@@ -204,7 +219,7 @@ def process_invoice_row(allocation, attrs, su_name, rate):
                     allocation,
                     [attributes.QUOTA_VOLUMES_GB, attributes.QUOTA_OBJECT_GB],
                     "OpenStack Storage",
-                    options['openstack_gb_rate'])
+                    openstack_storage_rate)
 
             for allocation in openshift_allocations:
                 allocation_str = f'{allocation.pk} of project "{allocation.project.title}"'
@@ -215,7 +230,7 @@ def process_invoice_row(allocation, attrs, su_name, rate):
                     allocation,
                     [attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, attributes.QUOTA_REQUESTS_STORAGE],
                     "OpenShift Storage",
-                    options['openshift_gb_rate']
+                    openshift_storage_rate
                 )
 
         if options['upload_to_s3']: