Add a session token to sqs config

airbytehq · Feb 21, 2024 · e57890f · e57890f
1 parent d520990
commit e57890f
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/airbyte-integrations/connectors/source-amazon-sqs/integration_tests/invalid_config.json b/airbyte-integrations/connectors/source-amazon-sqs/integration_tests/invalid_config.json
@@ -7,5 +7,6 @@
   "queue_url": "https://sqs.eu-west-1.amazonaws.com/840836244599/ab-airbyte-testing",
   "region": "eu-west-1",
   "access_key": "xxx",
-  "secret_key": "xxx"
+  "secret_key": "xxx",
+  "session_token": "xxx"
 }
diff --git a/airbyte-integrations/connectors/source-amazon-sqs/source_amazon_sqs/source.py b/airbyte-integrations/connectors/source-amazon-sqs/source_amazon_sqs/source.py
@@ -62,9 +62,12 @@ def check(self, logger: AirbyteLogger, config: json) -> AirbyteConnectionStatus:
             logger.debug("Amazon SQS Source Config Check - access_key (ends with): " + access_key[-1])
             secret_key = config["secret_key"]
             logger.debug("Amazon SQS Source Config Check - secret_key (ends with): " + secret_key[-1])
+            session_token = config.get("session_token")
+            if session_token:
+                logger.debug("Amazon SQS Source Config Check - secret_key (ends with): " + session_token[-1])
 
             logger.debug("Amazon SQS Source Config Check - Starting connection test ---")
-            session = boto3.Session(aws_access_key_id=access_key, aws_secret_access_key=secret_key, region_name=queue_region)
+            session = boto3.Session(aws_access_key_id=access_key, aws_secret_access_key=secret_key, aws_session_token=session_token, region_name=queue_region)
             sqs = session.resource("sqs")
             queue = sqs.Queue(url=queue_url)
             if hasattr(queue, "attributes"):

diff --git a/airbyte-integrations/connectors/source-amazon-sqs/source_amazon_sqs/spec.json b/airbyte-integrations/connectors/source-amazon-sqs/source_amazon_sqs/spec.json
@@ -107,6 +107,13 @@
         "examples": ["hu+qE5exxxxT6o/ZrKsxxxxxxBhxxXLexxxxxVKz"],
         "airbyte_secret": true,
         "order": 8
+      },
+      "session_token": {
+        "title": "AWS IAM Session Token",
+        "description": "The Temporary Session Token of the AWS IAM Role to use for pulling messages",
+        "type": "string",
+        "airbyte_secret": true,
+        "order": 9
       }
     }
   }