From 9b62d49689f5c32222bfb1b2014f29df219cfcd9 Mon Sep 17 00:00:00 2001
From: Mike Perez <thingee@gmail.com>
Date: Wed, 28 Aug 2024 18:43:38 -0700
Subject: [PATCH] Add s3 upload storage method

Signed-off-by: Mike Perez <thingee@gmail.com>
---
 chacra/controllers/binaries/archs.py | 26 +++++++++++++++++++++++---
 config/dev.py                        |  6 ++++++
 requirements.txt                     |  1 +
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/chacra/controllers/binaries/archs.py b/chacra/controllers/binaries/archs.py
index f2c63ec5..535ec15a 100644
--- a/chacra/controllers/binaries/archs.py
+++ b/chacra/controllers/binaries/archs.py
@@ -1,5 +1,7 @@
 import logging
 import os
+import boto3
+from botocore.exceptions import ClientError
 import pecan
 from pecan import response
 from pecan.secure import secure
@@ -89,7 +91,7 @@ def index_post(self):
                 if request.POST.get('force', False) is False:
                     error('/errors/invalid', 'resource already exists and "force" key was not used')
 
-        full_path = self.save_file(file_obj)
+        full_path, size = self.save_file(file_obj)
 
         if self.binary is None:
             path = full_path
@@ -102,7 +104,7 @@ def index_post(self):
             self.binary = Binary(
                 self.binary_name, self.project, arch=arch,
                 distro=distro, distro_version=distro_version,
-                ref=ref, sha1=sha1, path=path, size=os.path.getsize(path)
+                ref=ref, sha1=sha1, path=path, size=size
             )
         else:
             self.binary.path = full_path
@@ -110,6 +112,11 @@ def index_post(self):
         # check if this binary is interesting for other configured projects,
         # and if so, then mark those other repos so that they can be re-built
         self.mark_related_repos()
+
+        # Remove the local file after S3 upload
+        if pecan.conf.storage_method == 's3':
+            os.remove(full_path)
+
         return dict()
 
     def mark_related_repos(self):
@@ -175,8 +182,21 @@ def save_file(self, file_obj):
             for chunk in file_iterable:
                 f.write(chunk)
 
+        if pecan.conf.storage_method == 's3':
+            bucket = pecan.conf.bucket
+            object_name = os.path.basename(self.binary_name)
+
+            s3_client = boto3.client('s3')
+            try:
+                with open(destination, 'rb') as f:
+                    s3_client.upload_fileobj(f, bucket, object_name)
+            except ClientError as e:
+                error('/errors/error/', 'file object upload to S3 failed with error %s' % e)
+
+        size = os.path.getsize(destination)
+
         # return the full path to the saved object:
-        return destination
+        return destination, size
 
     @expose()
     def _lookup(self, name, *remainder):
diff --git a/config/dev.py b/config/dev.py
index 917556d1..6c5b0425 100644
--- a/config/dev.py
+++ b/config/dev.py
@@ -63,11 +63,17 @@
     'encoding':      'utf-8'
 }
 
+# Where to store the data. Options are 's3' or 'local'
+storage_method = 'local'
+
 # location for storing uploaded binaries
 binary_root = '%(confdir)s/public'
 repos_root = '%(confdir)s/repos'
 distributions_root = '%(confdir)s/distributions'
 
+# If storage method is s3, provide a bucket name
+bucket = ''
+
 # When True it will set the headers so that Nginx can serve the download
 # instead of Pecan.
 delegate_downloads = False
diff --git a/requirements.txt b/requirements.txt
index 42b61cde..63454df6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,4 +8,5 @@ alembic
 ipython
 python-statsd
 requests
+boto3
 importlib_metadata<=3.6; python_version<'3.8'