Merge pull request #54 from knikolla/arg_skip_primary

Add command line argument to upload or skip primary s3 location

src/openstack_billing_db/billing.py

@@ -1,4 +1,5 @@
 import csv
+import logging
 from datetime import datetime
 from dataclasses import dataclass
 from decimal import Decimal
@@ -10,6 +11,8 @@
 
 import boto3
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass()
 class Rates(object):
@@ -214,7 +217,8 @@ def generate_billing(start, end, output, rates,
                      coldfront_data_file=None,
                      invoice_month=None,
                      upload_to_s3=False,
-                     sql_dump_file=None):
+                     sql_dump_file=None,
+                     upload_to_primary_location=True):
 
     database = model.Database(start, sql_dump_file)
 
@@ -243,15 +247,18 @@ def generate_billing(start, end, output, rates,
             aws_secret_access_key=s3_secret,
         )
 
-        primary_location = (
-            f"Invoices/{invoice_month}/"
-            f"Service Invoices/NERC OpenStack {invoice_month}.csv"
-        )
-        s3.upload_file(output, Bucket=s3_bucket, Key=primary_location)
+        if upload_to_primary_location:
+            primary_location = (
+                f"Invoices/{invoice_month}/"
+                f"Service Invoices/NERC OpenStack {invoice_month}.csv"
+            )
+            s3.upload_file(output, Bucket=s3_bucket, Key=primary_location)
+            logger.info(f"Uploaded to {primary_location}.")
 
         timestamp = datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')
         secondary_location = (
             f"Invoices/{invoice_month}/"
             f"Archive/NERC OpenStack {invoice_month} {timestamp}.csv"
         )
         s3.upload_file(output, Bucket=s3_bucket, Key=secondary_location)
+        logger.info(f"Uploaded to {secondary_location}.")

src/openstack_billing_db/main.py

@@ -148,6 +148,13 @@ def main():
               " but can be configured through S3_OUTPUT_BUCKET and"
               " S3_OUTPUT_ENDPOINT_URL environment variables.")
     )
+    parser.add_argument(
+        "--upload-to-primary-location",
+        default=True,
+        type=bool,
+        help=("When uploading to S3, upload both to primary and"
+              " archive location, or just archive location.")
+    )
     parser.add_argument(
         "--output-file",
         default="/tmp/openstack_invoices.csv",
@@ -198,6 +205,7 @@ def main():
         invoice_month=args.invoice_month,
         upload_to_s3=args.upload_to_s3,
         sql_dump_file=dump_file,
+        upload_to_primary_location=args.upload_to_primary_location,
     )