Skip to content

Commit 0057720

Browse files
rohan-patil2rohan-patil2
authored
feat: add AWS STS AssumeRole support to AWS credential (#5731)
* feat: add AWS STS AssumeRole support to AWS credential * fix: applied reviewer suggestions for code hardening * fix: minor lint issues * fix: resolved lint warnings
1 parent ec767b1 commit 0057720

12 files changed

Lines changed: 310 additions & 204 deletions

File tree

‎packages/components/credentials/AWSCredential.credential.ts‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ class AWSApi implements INodeCredential {
1111
constructor() {
1212
this.label = 'AWS security credentials'
1313
this.name = 'awsApi'
14-
this.version = 1.0
14+
this.version = 1.1
1515
this.description =
1616
'Your <a target="_blank" href="https://docs.aws.amazon.com/IAM/latest/UserGuide/security-creds.html">AWS security credentials</a>. When unspecified, credentials will be sourced from the runtime environment according to the default AWS SDK behavior.'
1717
this.optional = true
@@ -39,6 +39,24 @@ class AWSApi implements INodeCredential {
3939
placeholder: '<AWS_SESSION_TOKEN>',
4040
description: 'The session key for your AWS account. This is only needed when you are using temporary credentials.',
4141
optional: true
42+
},
43+
{
44+
label: 'Role ARN',
45+
name: 'roleArn',
46+
type: 'string',
47+
placeholder: 'arn:aws:iam::123456789012:role/role-name',
48+
description:
49+
'The Amazon Resource Name (ARN) of the IAM role to assume. When provided, Flowise will use AWS STS AssumeRole to obtain temporary credentials. Leave empty to use static credentials directly.',
50+
optional: true
51+
},
52+
{
53+
label: 'External ID',
54+
name: 'externalId',
55+
type: 'string',
56+
placeholder: 'unique-external-id',
57+
description:
58+
'A unique identifier used for cross-account role assumption. Required when the role trust policy includes an sts:ExternalId condition.',
59+
optional: true
4260
}
4361
]
4462
}

‎packages/components/nodes/chatmodels/AWSBedrock/AWSChatBedrock.ts‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import { BaseCache } from '@langchain/core/caches'
22
import { ICommonObject, IMultiModalOption, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../src/Interface'
3-
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
3+
import { getBaseClasses } from '../../../src/utils'
44
import { getModels, getRegions, MODEL_TYPE } from '../../../src/modelLoader'
5+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
56
import { ChatBedrockConverseInput, ChatBedrockConverse } from '@langchain/aws'
67
import { BedrockChat } from './FlowiseAWSChatBedrock'
78

@@ -163,19 +164,12 @@ class AWSChatBedrock_ChatModels implements INode {
163164
* Bedrock's credential provider falls back to the AWS SDK to fetch
164165
* credentials from the running environment.
165166
* When specified, we override the default provider with configured values.
167+
* Supports STS AssumeRole when a Role ARN is configured in the credential.
166168
* @see https://github.com/aws/aws-sdk-js-v3/blob/main/packages/credential-provider-node/README.md
167169
*/
168-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
169-
if (credentialData && Object.keys(credentialData).length !== 0) {
170-
const credentialApiKey = getCredentialParam('awsKey', credentialData, nodeData)
171-
const credentialApiSecret = getCredentialParam('awsSecret', credentialData, nodeData)
172-
const credentialApiSession = getCredentialParam('awsSession', credentialData, nodeData)
173-
174-
obj.credentials = {
175-
accessKeyId: credentialApiKey,
176-
secretAccessKey: credentialApiSecret,
177-
sessionToken: credentialApiSession
178-
}
170+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, iRegion)
171+
if (credentialConfig.credentials) {
172+
obj.credentials = credentialConfig.credentials
179173
}
180174
if (cache) obj.cache = cache
181175

‎packages/components/nodes/documentloaders/S3Directory/S3Directory.ts‎

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,6 @@
11
import { ICommonObject, INode, INodeData, INodeOptionsValue, INodeOutputsValue, INodeParams } from '../../../src/Interface'
2-
import {
3-
getCredentialData,
4-
getCredentialParam,
5-
handleDocumentLoaderDocuments,
6-
handleDocumentLoaderMetadata,
7-
handleDocumentLoaderOutput
8-
} from '../../../src/utils'
2+
import { handleDocumentLoaderDocuments, handleDocumentLoaderMetadata, handleDocumentLoaderOutput } from '../../../src/utils'
3+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
94
import { S3Client, GetObjectCommand, S3ClientConfig, ListObjectsV2Command, ListObjectsV2Output } from '@aws-sdk/client-s3'
105
import { getRegions, MODEL_TYPE } from '../../../src/modelLoader'
116
import { Readable } from 'node:stream'
@@ -158,18 +153,9 @@ class S3_DocumentLoaders implements INode {
158153
const output = nodeData.outputs?.output as string
159154

160155
let credentials: S3ClientConfig['credentials'] | undefined
161-
162156
if (nodeData.credential) {
163-
const credentialData = await getCredentialData(nodeData.credential, options)
164-
const accessKeyId = getCredentialParam('awsKey', credentialData, nodeData)
165-
const secretAccessKey = getCredentialParam('awsSecret', credentialData, nodeData)
166-
167-
if (accessKeyId && secretAccessKey) {
168-
credentials = {
169-
accessKeyId,
170-
secretAccessKey
171-
}
172-
}
157+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
158+
credentials = credentialConfig.credentials
173159
}
174160

175161
let s3Config: S3ClientConfig = {

‎packages/components/nodes/documentloaders/S3File/S3File.ts‎

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,8 @@ import {
77
SkipInferTableTypes,
88
HiResModelName
99
} from '@langchain/community/document_loaders/fs/unstructured'
10-
import {
11-
getCredentialData,
12-
getCredentialParam,
13-
handleDocumentLoaderDocuments,
14-
handleDocumentLoaderMetadata,
15-
handleDocumentLoaderOutput
16-
} from '../../../src/utils'
10+
import { handleDocumentLoaderDocuments, handleDocumentLoaderMetadata, handleDocumentLoaderOutput } from '../../../src/utils'
11+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
1712
import { S3Client, GetObjectCommand, HeadObjectCommand, S3ClientConfig } from '@aws-sdk/client-s3'
1813
import { getRegions, MODEL_TYPE } from '../../../src/modelLoader'
1914
import { Readable } from 'node:stream'
@@ -581,18 +576,9 @@ class S3_DocumentLoaders implements INode {
581576
}
582577

583578
let credentials: S3ClientConfig['credentials'] | undefined
584-
585579
if (nodeData.credential) {
586-
const credentialData = await getCredentialData(nodeData.credential, options)
587-
const accessKeyId = getCredentialParam('awsKey', credentialData, nodeData)
588-
const secretAccessKey = getCredentialParam('awsSecret', credentialData, nodeData)
589-
590-
if (accessKeyId && secretAccessKey) {
591-
credentials = {
592-
accessKeyId,
593-
secretAccessKey
594-
}
595-
}
580+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
581+
credentials = credentialConfig.credentials
596582
}
597583

598584
const s3Config: S3ClientConfig = {

‎packages/components/nodes/embeddings/AWSBedrockEmbedding/AWSBedrockEmbedding.ts‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime'
22
import { BedrockEmbeddings, BedrockEmbeddingsParams } from '@langchain/community/embeddings/bedrock'
33
import { ICommonObject, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../src/Interface'
4-
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
4+
import { getBaseClasses } from '../../../src/utils'
5+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
56
import { MODEL_TYPE, getModels, getRegions } from '../../../src/modelLoader'
67

78
class AWSBedrockEmbedding_Embeddings implements INode {
@@ -129,17 +130,15 @@ class AWSBedrockEmbedding_Embeddings implements INode {
129130
region: iRegion
130131
}
131132

132-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
133-
if (credentialData && Object.keys(credentialData).length !== 0) {
134-
const credentialApiKey = getCredentialParam('awsKey', credentialData, nodeData)
135-
const credentialApiSecret = getCredentialParam('awsSecret', credentialData, nodeData)
136-
const credentialApiSession = getCredentialParam('awsSession', credentialData, nodeData)
137-
138-
obj.credentials = {
139-
accessKeyId: credentialApiKey,
140-
secretAccessKey: credentialApiSecret,
141-
sessionToken: credentialApiSession
142-
}
133+
/**
134+
* Long-term credentials specified in embedding configuration are optional.
135+
* Bedrock's credential provider falls back to the AWS SDK to fetch
136+
* credentials from the running environment.
137+
* Supports STS AssumeRole when a Role ARN is configured in the credential.
138+
*/
139+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, iRegion)
140+
if (credentialConfig.credentials) {
141+
obj.credentials = credentialConfig.credentials
143142
}
144143

145144
const client = new BedrockRuntimeClient({

‎packages/components/nodes/llms/AWSBedrock/AWSBedrock.ts‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@ import { Bedrock } from '@langchain/community/llms/bedrock'
22
import { BaseCache } from '@langchain/core/caches'
33
import { BaseLLMParams } from '@langchain/core/language_models/llms'
44
import { ICommonObject, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../src/Interface'
5-
import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
5+
import { getBaseClasses } from '../../../src/utils'
66
import { BaseBedrockInput } from '@langchain/community/dist/utils/bedrock'
77
import { getModels, getRegions, MODEL_TYPE } from '../../../src/modelLoader'
8+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
89

910
/**
1011
* @author Michael Connor <mlconnor@yahoo.com>
@@ -116,19 +117,12 @@ class AWSBedrock_LLMs implements INode {
116117
* Bedrock's credential provider falls back to the AWS SDK to fetch
117118
* credentials from the running environment.
118119
* When specified, we override the default provider with configured values.
120+
* Supports STS AssumeRole when a Role ARN is configured in the credential.
119121
* @see https://github.com/aws/aws-sdk-js-v3/blob/main/packages/credential-provider-node/README.md
120122
*/
121-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
122-
if (credentialData && Object.keys(credentialData).length !== 0) {
123-
const credentialApiKey = getCredentialParam('awsKey', credentialData, nodeData)
124-
const credentialApiSecret = getCredentialParam('awsSecret', credentialData, nodeData)
125-
const credentialApiSession = getCredentialParam('awsSession', credentialData, nodeData)
126-
127-
obj.credentials = {
128-
accessKeyId: credentialApiKey,
129-
secretAccessKey: credentialApiSecret,
130-
sessionToken: credentialApiSession
131-
}
123+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, iRegion)
124+
if (credentialConfig.credentials) {
125+
obj.credentials = credentialConfig.credentials
132126
}
133127
if (cache) obj.cache = cache
134128

‎packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.ts‎

Lines changed: 5 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { AmazonKnowledgeBaseRetriever } from '@langchain/aws'
22
import { ICommonObject, INode, INodeData, INodeParams, INodeOptionsValue } from '../../../src/Interface'
3-
import { getCredentialData, getCredentialParam } from '../../../src/utils'
3+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
44
import { RetrievalFilter } from '@aws-sdk/client-bedrock-agent-runtime'
55
import { MODEL_TYPE, getRegions } from '../../../src/modelLoader'
66

@@ -108,29 +108,15 @@ class AWSBedrockKBRetriever_Retrievers implements INode {
108108
const topK = nodeData.inputs?.topK as number
109109
const overrideSearchType = (nodeData.inputs?.searchType != '' ? nodeData.inputs?.searchType : undefined) as 'HYBRID' | 'SEMANTIC'
110110
const filter = (nodeData.inputs?.filter != '' ? JSON.parse(nodeData.inputs?.filter) : undefined) as RetrievalFilter
111-
let credentialApiKey = ''
112-
let credentialApiSecret = ''
113-
let credentialApiSession = ''
114-
115-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
116-
if (credentialData && Object.keys(credentialData).length !== 0) {
117-
credentialApiKey = getCredentialParam('awsKey', credentialData, nodeData)
118-
credentialApiSecret = getCredentialParam('awsSecret', credentialData, nodeData)
119-
credentialApiSession = getCredentialParam('awsSession', credentialData, nodeData)
120-
}
121-
111+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
122112
const retriever = new AmazonKnowledgeBaseRetriever({
123-
topK: topK,
113+
topK,
124114
knowledgeBaseId: knoledgeBaseID,
125-
region: region,
115+
region,
126116
filter,
127117
overrideSearchType,
128118
clientOptions: {
129-
credentials: {
130-
accessKeyId: credentialApiKey,
131-
secretAccessKey: credentialApiSecret,
132-
sessionToken: credentialApiSession
133-
}
119+
...(credentialConfig.credentials && { credentials: credentialConfig.credentials })
134120
}
135121
})
136122

‎packages/components/nodes/tools/AWSDynamoDBKVStorage/AWSDynamoDBKVStorage.ts‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,18 @@ const KEY_SEPARATOR = '#'
1717
const MAX_KEY_LENGTH = 2048 // DynamoDB limit for partition key
1818

1919
// Helper function to create DynamoDB client
20-
function createDynamoDBClient(credentials: AWSCredentials, region: string): DynamoDBClient {
21-
return new DynamoDBClient({
22-
region,
23-
credentials: {
20+
function createDynamoDBClient(credentials: AWSCredentials | undefined, region: string): DynamoDBClient {
21+
const config: { region: string; credentials?: { accessKeyId: string; secretAccessKey: string; sessionToken?: string } } = { region }
22+
23+
if (credentials) {
24+
config.credentials = {
2425
accessKeyId: credentials.accessKeyId,
2526
secretAccessKey: credentials.secretAccessKey,
2627
...(credentials.sessionToken && { sessionToken: credentials.sessionToken })
2728
}
28-
})
29+
}
30+
31+
return new DynamoDBClient(config)
2932
}
3033

3134
// Helper function to build full key with optional prefix
@@ -246,11 +249,11 @@ class AWSDynamoDBKVStorage_Tools implements INode {
246249
]
247250
}
248251

249-
loadMethods: Record<string, (nodeData: INodeData, options?: ICommonObject) => Promise<INodeOptionsValue[]>> = {
250-
listTables: async (nodeData: INodeData, options?: ICommonObject): Promise<INodeOptionsValue[]> => {
252+
loadMethods: Record<string, (_nodeData: INodeData, _options?: ICommonObject) => Promise<INodeOptionsValue[]>> = {
253+
listTables: async (_nodeData: INodeData, _options?: ICommonObject): Promise<INodeOptionsValue[]> => {
251254
try {
252-
const credentials = await getAWSCredentials(nodeData, options ?? {})
253-
const region = (nodeData.inputs?.region as string) || DEFAULT_AWS_REGION
255+
const credentials = await getAWSCredentials(_nodeData, _options ?? {})
256+
const region = (_nodeData.inputs?.region as string) || DEFAULT_AWS_REGION
254257
const dynamoClient = createDynamoDBClient(credentials, region)
255258

256259
const listCommand = new ListTablesCommand({})

‎packages/components/nodes/vectorstores/Kendra/Kendra.ts‎

Lines changed: 15 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
import { flatten } from 'lodash'
22
import { AmazonKendraRetriever } from '@langchain/aws'
3-
import { KendraClient, BatchPutDocumentCommand, BatchDeleteDocumentCommand } from '@aws-sdk/client-kendra'
3+
import { KendraClient, BatchPutDocumentCommand, BatchDeleteDocumentCommand, KendraClientConfig } from '@aws-sdk/client-kendra'
44
import { Document } from '@langchain/core/documents'
55
import { ICommonObject, INode, INodeData, INodeOptionsValue, INodeOutputsValue, INodeParams, IndexingResult } from '../../../src/Interface'
6-
import { FLOWISE_CHATID, getCredentialData, getCredentialParam, parseJsonBody } from '../../../src/utils'
6+
import { FLOWISE_CHATID, parseJsonBody } from '../../../src/utils'
7+
import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
78
import { howToUseFileUpload } from '../VectorStoreUtils'
89
import { MODEL_TYPE, getRegions } from '../../../src/modelLoader'
910

@@ -119,23 +120,11 @@ class Kendra_VectorStores implements INode {
119120
const docs = nodeData.inputs?.document as Document[]
120121
const isFileUploadEnabled = nodeData.inputs?.fileUpload as boolean
121122

122-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
123-
let clientConfig: any = { region }
124-
125-
if (credentialData && Object.keys(credentialData).length !== 0) {
126-
const accessKeyId = getCredentialParam('awsKey', credentialData, nodeData)
127-
const secretAccessKey = getCredentialParam('awsSecret', credentialData, nodeData)
128-
const sessionToken = getCredentialParam('awsSession', credentialData, nodeData)
129-
130-
if (accessKeyId && secretAccessKey) {
131-
clientConfig.credentials = {
132-
accessKeyId,
133-
secretAccessKey,
134-
...(sessionToken && { sessionToken })
135-
}
136-
}
123+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
124+
let clientConfig: KendraClientConfig = { region }
125+
if (credentialConfig.credentials) {
126+
clientConfig.credentials = credentialConfig.credentials
137127
}
138-
139128
const client = new KendraClient(clientConfig)
140129

141130
const flattenDocs = docs && docs.length ? flatten(docs) : []
@@ -192,23 +181,11 @@ class Kendra_VectorStores implements INode {
192181
const indexId = nodeData.inputs?.indexId as string
193182
const region = nodeData.inputs?.region as string
194183

195-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
196-
let clientConfig: any = { region }
197-
198-
if (credentialData && Object.keys(credentialData).length !== 0) {
199-
const accessKeyId = getCredentialParam('awsKey', credentialData, nodeData)
200-
const secretAccessKey = getCredentialParam('awsSecret', credentialData, nodeData)
201-
const sessionToken = getCredentialParam('awsSession', credentialData, nodeData)
202-
203-
if (accessKeyId && secretAccessKey) {
204-
clientConfig.credentials = {
205-
accessKeyId,
206-
secretAccessKey,
207-
...(sessionToken && { sessionToken })
208-
}
209-
}
184+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
185+
let clientConfig: KendraClientConfig = { region }
186+
if (credentialConfig.credentials) {
187+
clientConfig.credentials = credentialConfig.credentials
210188
}
211-
212189
const client = new KendraClient(clientConfig)
213190

214191
try {
@@ -235,15 +212,11 @@ class Kendra_VectorStores implements INode {
235212
const attributeFilter = nodeData.inputs?.attributeFilter
236213
const isFileUploadEnabled = nodeData.inputs?.fileUpload as boolean
237214

238-
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
239-
let clientOptions: any = {}
215+
const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
216+
let clientOptions: Partial<KendraClientConfig> = {}
240217

241-
if (credentialData && Object.keys(credentialData).length !== 0) {
242-
clientOptions.credentials = {
243-
accessKeyId: getCredentialParam('awsKey', credentialData, nodeData),
244-
secretAccessKey: getCredentialParam('awsSecret', credentialData, nodeData),
245-
sessionToken: getCredentialParam('awsSession', credentialData, nodeData)
246-
}
218+
if (credentialConfig.credentials) {
219+
clientOptions.credentials = credentialConfig.credentials
247220
}
248221

249222
let filter = undefined

0 commit comments

Comments
 (0)