From c6a8f1406c928de70bb1192329d62a7dc3f784e6 Mon Sep 17 00:00:00 2001
From: Mauro Cortellazzi <mauro.cortellazzi@radicalbit.io>
Date: Tue, 25 Jun 2024 17:58:34 +0200
Subject: [PATCH] feat(sdk): added ability to set S3 endpoint url if needed
 (#34)

---
 sdk/radicalbit_platform_sdk/apis/model.py     | 36 +++++++++++++++++++
 .../models/aws_credentials.py                 |  3 ++
 2 files changed, 39 insertions(+)

diff --git a/sdk/radicalbit_platform_sdk/apis/model.py b/sdk/radicalbit_platform_sdk/apis/model.py
index 739aaf3f..271d5489 100644
--- a/sdk/radicalbit_platform_sdk/apis/model.py
+++ b/sdk/radicalbit_platform_sdk/apis/model.py
@@ -181,6 +181,15 @@ def load_reference_dataset(
                         if aws_credentials is None
                         else aws_credentials.default_region
                     ),
+                    endpoint_url=(
+                        None
+                        if aws_credentials is None
+                        else (
+                            None
+                            if aws_credentials.endpoint_url is None
+                            else aws_credentials.endpoint_url
+                        )
+                    ),
                 )
 
                 s3_client.upload_file(
@@ -237,6 +246,15 @@ def bind_reference_dataset(
                 region_name=(
                     None if aws_credentials is None else aws_credentials.default_region
                 ),
+                endpoint_url=(
+                    None
+                    if aws_credentials is None
+                    else (
+                        None
+                        if aws_credentials.endpoint_url is None
+                        else aws_credentials.endpoint_url
+                    )
+                ),
             )
 
             chunks_iterator = s3_client.get_object(
@@ -314,6 +332,15 @@ def load_current_dataset(
                         if aws_credentials is None
                         else aws_credentials.default_region
                     ),
+                    endpoint_url=(
+                        None
+                        if aws_credentials is None
+                        else (
+                            None
+                            if aws_credentials.endpoint_url is None
+                            else aws_credentials.endpoint_url
+                        )
+                    ),
                 )
 
                 s3_client.upload_file(
@@ -372,6 +399,15 @@ def bind_current_dataset(
                 region_name=(
                     None if aws_credentials is None else aws_credentials.default_region
                 ),
+                endpoint_url=(
+                    None
+                    if aws_credentials is None
+                    else (
+                        None
+                        if aws_credentials.endpoint_url is None
+                        else aws_credentials.endpoint_url
+                    )
+                ),
             )
 
             chunks_iterator = s3_client.get_object(
diff --git a/sdk/radicalbit_platform_sdk/models/aws_credentials.py b/sdk/radicalbit_platform_sdk/models/aws_credentials.py
index a629b78f..eaf34956 100644
--- a/sdk/radicalbit_platform_sdk/models/aws_credentials.py
+++ b/sdk/radicalbit_platform_sdk/models/aws_credentials.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from pydantic import BaseModel
 
 
@@ -5,3 +7,4 @@ class AwsCredentials(BaseModel):
     access_key_id: str
     secret_access_key: str
     default_region: str
+    endpoint_url: Optional[str]