fix(CE): Add STS credentials for AWS S3 source connector (#224)

* fix(CE): Add STS credentials for AWS S3 source connector --------- Co-authored-by: pabss-ai2 <[email protected]> Co-authored-by: Pablo Rivera Bengoechea <[email protected]>
Multiwoven · Jul 2, 2024 · 64b66f1 · 64b66f1
1 parent e8542bd
commit 64b66f1
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 12 deletions.
diff --git a/integrations/Gemfile.lock b/integrations/Gemfile.lock
@@ -7,7 +7,7 @@ GIT
 PATH
   remote: .
   specs:
-    multiwoven-integrations (0.3.3)
+    multiwoven-integrations (0.3.4)
       activesupport
       async-websocket
       aws-sdk-athena

diff --git a/integrations/lib/multiwoven/integrations/core/constants.rb b/integrations/lib/multiwoven/integrations/core/constants.rb
@@ -34,6 +34,9 @@ module Constants
       AIRTABLE_BASES_ENDPOINT = "https://api.airtable.com/v0/meta/bases"
       AIRTABLE_GET_BASE_SCHEMA_ENDPOINT = "https://api.airtable.com/v0/meta/bases/{baseId}/tables"
 
+      AWS_ACCESS_KEY_ID = ENV["AWS_ACCESS_KEY_ID"]
+      AWS_SECRET_ACCESS_KEY = ENV["AWS_SECRET_ACCESS_KEY"]
+
       # HTTP
       HTTP_GET = "GET"
       HTTP_POST = "POST"

diff --git a/integrations/lib/multiwoven/integrations/rollout.rb b/integrations/lib/multiwoven/integrations/rollout.rb
@@ -2,7 +2,7 @@
 
 module Multiwoven
   module Integrations
-    VERSION = "0.3.3"
+    VERSION = "0.3.4"
 
     ENABLED_SOURCES = %w[
       Snowflake

diff --git a/integrations/lib/multiwoven/integrations/source/amazon_s3/client.rb b/integrations/lib/multiwoven/integrations/source/amazon_s3/client.rb
@@ -50,15 +50,17 @@ def read(sync_config)
       private
 
       def get_auth_data(connection_config)
-        session = @session_name
+        session = @session_name.gsub(/\s+/, "-")
         @session_name = ""
         if connection_config[:auth_type] == "user"
           Aws::Credentials.new(connection_config[:access_id], connection_config[:secret_access])
         elsif connection_config[:auth_type] == "role"
-          sts_client = Aws::STS::Client.new(region: connection_config[:region])
+          credentials = Aws::Credentials.new(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)
+          sts_client = Aws::STS::Client.new(region: connection_config[:region], credentials: credentials)
           resp = sts_client.assume_role({
                                           role_arn: connection_config[:arn],
-                                          role_session_name: session
+                                          role_session_name: session,
+                                          external_id: connection_config[:external_id]
                                         })
           Aws::Credentials.new(
             resp.credentials.access_key_id,

diff --git a/integrations/lib/multiwoven/integrations/source/amazon_s3/config/spec.json b/integrations/lib/multiwoven/integrations/source/amazon_s3/config/spec.json
@@ -29,6 +29,7 @@
         "region",
         "bucket",
         "arn",
+        "external_id",
         "file_type"
       ]
     },
@@ -53,16 +54,22 @@
         "title": "IAM Role ARN",
         "order": 1
       },
+      "external_id": {
+        "type": "string",
+        "title": "External Id",
+        "description": "Unique ID that allows handshake between AWS accounts.",
+        "order": 2
+      },
       "access_id": {
         "type": "string",
         "title": "Access Id",
-        "order": 2
+        "order": 3
       },
       "secret_access": {
         "type": "string",
         "title": "Secret Access",
         "multiwoven_secret": true,
-        "order": 3
+        "order": 4
       },
       "region": {
         "description": "AWS region",
@@ -71,13 +78,13 @@
         ],
         "type": "string",
         "title": "Region",
-        "order": 4
+        "order": 5
       },
       "bucket": {
         "description": "Bucket Name",
         "type": "string",
         "title": "Bucket",
-        "order": 5
+        "order": 6
       },
       "path": {
         "description": "Path to csv or parquet files",
@@ -86,7 +93,7 @@
         ],
         "type": "string",
         "title": "Path",
-        "order": 6
+        "order": 7
       },
       "file_type": {
         "description": "The type of file to read",
@@ -96,7 +103,7 @@
           "csv",
           "parquet"
         ],
-        "order": 7
+        "order": 8
       }
     }
   }

diff --git a/integrations/spec/multiwoven/integrations/source/amazons3/client_spec.rb b/integrations/spec/multiwoven/integrations/source/amazons3/client_spec.rb
@@ -17,7 +17,8 @@
           "access_id": "accessid",
           "secret_access": "secretaccess",
           "file_type": "type",
-          "arn": ""
+          "arn": "",
+          "external_id": ""
         }
       },
       "destination": {
@@ -76,6 +77,7 @@
         sync_config[:source][:connection_specification][:acess_id] = ""
         sync_config[:source][:connection_specification][:secret_access] = ""
         sync_config[:source][:connection_specification][:arn] = "aimrole/arn"
+        sync_config[:source][:connection_specification][:external_id] = "aws-external-id-trust-relationship"
         allow_any_instance_of(Multiwoven::Integrations::Source::AmazonS3::Client).to receive(:get_auth_data).and_return(auth_data)
         allow_any_instance_of(Multiwoven::Integrations::Source::AmazonS3::Client).to receive(:get_results).and_return([{ Id: "1" }, { Id: "2" }])
         message = client.check_connection(sync_config[:source][:connection_specification])
@@ -127,6 +129,7 @@
       sync_config[:source][:connection_specification][:acess_id] = ""
       sync_config[:source][:connection_specification][:secret_access] = ""
       sync_config[:source][:connection_specification][:arn] = "aimrole/arn"
+      sync_config[:source][:connection_specification][:external_id] = "aws-external-id-trust-relationship"
       s_config = Multiwoven::Integrations::Protocol::SyncConfig.from_json(sync_config.to_json)
       stub_request(:post, "https://sts.us-east-1.amazonaws.com/").to_return(status: 200, body: "", headers: {})
       allow(client).to receive(:get_auth_data).and_return(auth_data)
@@ -143,6 +146,7 @@
       sync_config[:source][:connection_specification][:acess_id] = ""
       sync_config[:source][:connection_specification][:secret_access] = ""
       sync_config[:source][:connection_specification][:arn] = "aimrole/arn"
+      sync_config[:source][:connection_specification][:external_id] = "aws-external-id-trust-relationship"
       s_config = Multiwoven::Integrations::Protocol::SyncConfig.from_json(sync_config.to_json)
       s_config.limit = 100
       s_config.offset = 1
@@ -197,6 +201,7 @@
       sync_config[:source][:connection_specification][:acess_id] = ""
       sync_config[:source][:connection_specification][:secret_access] = ""
       sync_config[:source][:connection_specification][:arn] = "aimrole/arn"
+      sync_config[:source][:connection_specification][:external_id] = "aws-external-id-trust-relationship"
       connection_config = sync_config[:source][:connection_specification]
       full_path = "s3://#{connection_config[:bucket]}/#{connection_config[:path]}*.#{connection_config[:file_type]}"
       allow(client).to receive(:get_auth_data).and_return(auth_data)