Merge pull request #12 from ppvolto/env-s3-endpoint-url

[S3Client] Add S3_ENTPOINT_URL #8
TemryL · Jun 15, 2024 · 9d89629 · 9d89629
2 parents 56b8e38 + 61fdd19
commit 9d89629
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/.env b/.env
@@ -3,4 +3,7 @@ S3_ACCESS_KEY = "replace with your access key"
 S3_SECRET_KEY = "replace with your secret key"
 S3_BUCKET_NAME = "replace with your bucket name"
 S3_INPUT_DIR = "replace with your S3 input dir"
-S3_OUTPUT_DIR = "replace with your S3 output dir"
+S3_OUTPUT_DIR = "replace with your S3 output dir"
+
+# Optional Enviroment Variables
+#S3_ENDPOINT_URL="replace with your S3 Endoint Url"
diff --git a/README.md b/README.md
@@ -22,6 +22,9 @@ S3_INPUT_DIR = "..."
 S3_OUTPUT_DIR = "..."
 ```
 
+### Optional S3 Config Variables
+- ```S3_ENDPOINT_URL``` allows the useage of a AWS Private Link or Other S3 Compatible Storage Solutions
+
 # Available Features
 ComfyUI nodes to:
 - [x] standalone download/upload file from/to Amazon S3

diff --git a/src/client_s3.py b/src/client_s3.py
@@ -8,11 +8,12 @@
 
 
 class S3:
-    def __init__(self, region, access_key, secret_key, bucket_name):
+    def __init__(self, region, access_key, secret_key, bucket_name, endpoint_url):
         self.region = region
         self.access_key = access_key
         self.secret_key = secret_key
         self.bucket_name = bucket_name
+        self.endpoint_url = endpoint_url
         self.s3_client = self.get_client()
         self.input_dir = os.getenv("S3_INPUT_DIR")
         self.output_dir = os.getenv("S3_OUTPUT_DIR")
@@ -31,7 +32,8 @@ def get_client(self):
                 service_name='s3',
                 region_name=self.region,
                 aws_access_key_id=self.access_key,
-                aws_secret_access_key=self.secret_key
+                aws_secret_access_key=self.secret_key,
+                endpoint_url=self.endpoint_url
             )
             return s3
         except Exception as e:
@@ -141,9 +143,10 @@ def get_s3_instance():
             region=os.getenv("S3_REGION"),
             access_key=os.getenv("S3_ACCESS_KEY"),
             secret_key=os.getenv("S3_SECRET_KEY"),
-            bucket_name=os.getenv("S3_BUCKET_NAME")
+            bucket_name=os.getenv("S3_BUCKET_NAME"),
+            endpoint_url=os.getenv("S3_ENDPOINT_URL")
         )
         return s3_instance
     except Exception as e:
         err = f"Failed to create S3 instance: {e} Please check your environment variables."
-        logger.error(err)
+        logger.error(err)