feat(storage): enables location credentials provider (#13605)

* feat: add location credentials provider * chore: add unit tests * chore: address feedback * chore: add locationCredentialsOption to copy * chore: remove casting types * chore: assert idenitity id * chore: avoid export common options interface * chore: address feedback * chore: fix test * chore: address feedback * address feedback * chore: clean-up types * chore: add test
aws-amplify · Jul 22, 2024 · c5464ac · c5464ac
1 parent 5d5bea1
commit c5464ac
Show file tree

Hide file tree

Showing 16 changed files with 235 additions and 26 deletions.
diff --git a/packages/storage/__tests__/providers/s3/apis/utils/resolveS3ConfigAndInput.test.ts b/packages/storage/__tests__/providers/s3/apis/utils/resolveS3ConfigAndInput.test.ts
@@ -9,6 +9,11 @@ import {
 	StorageValidationErrorCode,
 	validationErrorMap,
 } from '../../../../../src/errors/types/validation';
+import {
+	CallbackPathStorageInput,
+	DeprecatedStorageInput,
+} from '../../../../../src/providers/s3/utils/resolveS3ConfigAndInput';
+import { INVALID_STORAGE_INPUT } from '../../../../../src/errors/constants';
 
 jest.mock('@aws-amplify/core', () => ({
 	ConsoleLogger: jest.fn(),
@@ -76,13 +81,11 @@ describe('resolveS3ConfigAndInput', () => {
 		}
 	});
 
-	it('should throw if identityId is not available', async () => {
+	it('should not throw if identityId is not available', async () => {
 		mockFetchAuthSession.mockResolvedValueOnce({
 			credentials,
 		});
-		await expect(resolveS3ConfigAndInput(Amplify, {})).rejects.toMatchObject(
-			validationErrorMap[StorageValidationErrorCode.NoIdentityId],
-		);
+		expect(async () => resolveS3ConfigAndInput(Amplify, {})).not.toThrow();
 	});
 
 	it('should resolve bucket from S3 config', async () => {
@@ -179,7 +182,7 @@ describe('resolveS3ConfigAndInput', () => {
 	it('should resolve prefix with given access level', async () => {
 		mockDefaultResolvePrefix.mockResolvedValueOnce('prefix');
 		const { keyPrefix } = await resolveS3ConfigAndInput(Amplify, {
-			accessLevel: 'someLevel' as any,
+			options: { accessLevel: 'someLevel' as any },
 		});
 		expect(mockDefaultResolvePrefix).toHaveBeenCalledWith({
 			accessLevel: 'someLevel',
@@ -214,4 +217,91 @@ describe('resolveS3ConfigAndInput', () => {
 		});
 		expect(keyPrefix).toEqual('prefix');
 	});
+
+	describe('with locationCredentialsProvider', () => {
+		const mockLocationCredentialsProvider = jest
+			.fn()
+			.mockReturnValue({ credentials });
+		it('should resolve credentials without Amplify singleton', async () => {
+			mockGetConfig.mockReturnValue({
+				Storage: {
+					S3: {
+						bucket,
+						region,
+					},
+				},
+			});
+			const { s3Config } = await resolveS3ConfigAndInput(Amplify, {
+				options: {
+					locationCredentialsProvider: mockLocationCredentialsProvider,
+				},
+			});
+
+			if (typeof s3Config.credentials === 'function') {
+				const result = await s3Config.credentials();
+				expect(mockLocationCredentialsProvider).toHaveBeenCalled();
+				expect(result).toEqual(credentials);
+			} else {
+				throw new Error('Expect credentials to be a function');
+			}
+		});
+
+		it('should not throw when path is pass as a string', async () => {
+			const { s3Config } = await resolveS3ConfigAndInput(Amplify, {
+				path: 'my-path',
+				options: {
+					locationCredentialsProvider: mockLocationCredentialsProvider,
+				},
+			});
+
+			if (typeof s3Config.credentials === 'function') {
+				const result = await s3Config.credentials();
+				expect(mockLocationCredentialsProvider).toHaveBeenCalled();
+				expect(result).toEqual(credentials);
+			} else {
+				throw new Error('Expect credentials to be a function');
+			}
+		});
+
+		describe('with deprecated or callback paths as inputs', () => {
+			const key = 'mock-value';
+			const prefix = 'mock-value';
+			const path = () => 'path';
+			const deprecatedInputs: DeprecatedStorageInput[] = [
+				{ prefix },
+				{ key },
+				{
+					source: { key },
+					destination: { key },
+				},
+			];
+			const callbackPathInputs: CallbackPathStorageInput[] = [
+				{ path },
+				{
+					destination: { path },
+					source: { path },
+				},
+			];
+
+			const testCases = [...deprecatedInputs, ...callbackPathInputs];
+
+			it.each(testCases)('should throw when input is %s', async input => {
+				const { s3Config } = await resolveS3ConfigAndInput(Amplify, {
+					...input,
+					options: {
+						locationCredentialsProvider: mockLocationCredentialsProvider,
+					},
+				});
+				if (typeof s3Config.credentials === 'function') {
+					await expect(s3Config.credentials()).rejects.toThrow(
+						expect.objectContaining({
+							name: INVALID_STORAGE_INPUT,
+						}),
+					);
+				} else {
+					throw new Error('Expect credentials to be a function');
+				}
+			});
+		});
+	});
 });
diff --git a/packages/storage/src/errors/constants.ts b/packages/storage/src/errors/constants.ts
@@ -0,0 +1,4 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+export const INVALID_STORAGE_INPUT = 'InvalidStorageInput';
diff --git a/packages/storage/src/providers/s3/apis/downloadData.ts b/packages/storage/src/providers/s3/apis/downloadData.ts
@@ -115,7 +115,7 @@ const downloadDataJob =
 	> => {
 		const { options: downloadDataOptions } = downloadDataInput;
 		const { bucket, keyPrefix, s3Config, identityId } =
-			await resolveS3ConfigAndInput(Amplify, downloadDataOptions);
+			await resolveS3ConfigAndInput(Amplify, downloadDataInput);
 		const { inputType, objectKey } = validateStorageOperationInput(
 			downloadDataInput,
 			identityId,

diff --git a/packages/storage/src/providers/s3/apis/internal/copy.ts b/packages/storage/src/providers/s3/apis/internal/copy.ts
@@ -40,8 +40,10 @@ const copyWithPath = async (
 	input: CopyWithPathInput,
 ): Promise<CopyWithPathOutput> => {
 	const { source, destination } = input;
-	const { s3Config, bucket, identityId } =
-		await resolveS3ConfigAndInput(amplify);
+	const { s3Config, bucket, identityId } = await resolveS3ConfigAndInput(
+		amplify,
+		input,
+	);
 
 	assertValidationError(!!source.path, StorageValidationErrorCode.NoSourcePath);
 	assertValidationError(
@@ -92,10 +94,13 @@ export const copyWithKey = async (
 		s3Config,
 		bucket,
 		keyPrefix: sourceKeyPrefix,
-	} = await resolveS3ConfigAndInput(amplify, input.source);
+	} = await resolveS3ConfigAndInput(amplify, {
+		...input,
+		options: input.source,
+	});
 	const { keyPrefix: destinationKeyPrefix } = await resolveS3ConfigAndInput(
 		amplify,
-		input.destination,
+		{ ...input, options: input.destination },
 	); // resolveS3ConfigAndInput does not make extra API calls or storage access if called repeatedly.
 
 	// TODO(ashwinkumar6) V6-logger: warn `You may copy files from another user if the source level is "protected", currently it's ${srcLevel}`

diff --git a/packages/storage/src/providers/s3/apis/internal/getProperties.ts b/packages/storage/src/providers/s3/apis/internal/getProperties.ts
@@ -24,9 +24,8 @@ export const getProperties = async (
 	input: GetPropertiesInput | GetPropertiesWithPathInput,
 	action?: StorageAction,
 ): Promise<GetPropertiesOutput | GetPropertiesWithPathOutput> => {
-	const { options: getPropertiesOptions } = input;
 	const { s3Config, bucket, keyPrefix, identityId } =
-		await resolveS3ConfigAndInput(amplify, getPropertiesOptions);
+		await resolveS3ConfigAndInput(amplify, input);
 	const { inputType, objectKey } = validateStorageOperationInput(
 		input,
 		identityId,

diff --git a/packages/storage/src/providers/s3/apis/internal/getUrl.ts b/packages/storage/src/providers/s3/apis/internal/getUrl.ts
@@ -31,7 +31,7 @@ export const getUrl = async (
 ): Promise<GetUrlOutput | GetUrlWithPathOutput> => {
 	const { options: getUrlOptions } = input;
 	const { s3Config, keyPrefix, bucket, identityId } =
-		await resolveS3ConfigAndInput(amplify, getUrlOptions);
+		await resolveS3ConfigAndInput(amplify, input);
 	const { inputType, objectKey } = validateStorageOperationInput(
 		input,
 		identityId,

diff --git a/packages/storage/src/providers/s3/apis/internal/list.ts b/packages/storage/src/providers/s3/apis/internal/list.ts
@@ -58,7 +58,7 @@ export const list = async (
 		bucket,
 		keyPrefix: generatedPrefix,
 		identityId,
-	} = await resolveS3ConfigAndInput(amplify, options);
+	} = await resolveS3ConfigAndInput(amplify, input);
 
 	const { inputType, objectKey } = validateStorageOperationInputWithPrefix(
 		input,

diff --git a/packages/storage/src/providers/s3/apis/internal/remove.ts b/packages/storage/src/providers/s3/apis/internal/remove.ts
@@ -23,9 +23,8 @@ export const remove = async (
 	amplify: AmplifyClassV6,
 	input: RemoveInput | RemoveWithPathInput,
 ): Promise<RemoveOutput | RemoveWithPathOutput> => {
-	const { options = {} } = input ?? {};
 	const { s3Config, keyPrefix, bucket, identityId } =
-		await resolveS3ConfigAndInput(amplify, options);
+		await resolveS3ConfigAndInput(amplify, input);
 
 	const { inputType, objectKey } = validateStorageOperationInput(
 		input,

diff --git a/packages/storage/src/providers/s3/apis/uploadData/multipart/uploadHandlers.ts b/packages/storage/src/providers/s3/apis/uploadData/multipart/uploadHandlers.ts
@@ -73,7 +73,7 @@ export const getMultipartUploadHandlers = (
 		const { options: uploadDataOptions, data } = uploadDataInput;
 		const resolvedS3Options = await resolveS3ConfigAndInput(
 			Amplify,
-			uploadDataOptions,
+			uploadDataInput,
 		);
 
 		abortController = new AbortController();

diff --git a/packages/storage/src/providers/s3/apis/uploadData/putObjectJob.ts b/packages/storage/src/providers/s3/apis/uploadData/putObjectJob.ts
@@ -29,7 +29,7 @@ export const putObjectJob =
 	async (): Promise<ItemWithKey | ItemWithPath> => {
 		const { options: uploadDataOptions, data } = uploadDataInput;
 		const { bucket, keyPrefix, s3Config, isObjectLockEnabled, identityId } =
-			await resolveS3ConfigAndInput(Amplify, uploadDataOptions);
+			await resolveS3ConfigAndInput(Amplify, uploadDataInput);
 		const { inputType, objectKey } = validateStorageOperationInput(
 			uploadDataInput,
 			identityId,

diff --git a/packages/storage/src/providers/s3/types/inputs.ts b/packages/storage/src/providers/s3/types/inputs.ts
@@ -35,6 +35,8 @@ import {
 	UploadDataOptionsWithPath,
 } from '../types';
 
+import { LocationCredentialsProvider } from './options';
+
 // TODO: support use accelerate endpoint option
 /**
  * @deprecated Use {@link CopyWithPathInput} instead.
@@ -47,7 +49,9 @@ export type CopyInput = StorageCopyInputWithKey<
 /**
  * Input type with path for S3 copy API.
  */
-export type CopyWithPathInput = StorageCopyInputWithPath;
+export type CopyWithPathInput = StorageCopyInputWithPath<{
+	locationCredentialsProvider?: LocationCredentialsProvider;
+}>;
 
 /**
  * @deprecated Use {@link GetPropertiesWithPathInput} instead.

diff --git a/packages/storage/src/providers/s3/utils/resolveIdentityId.ts b/packages/storage/src/providers/s3/utils/resolveIdentityId.ts
@@ -0,0 +1,11 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { StorageValidationErrorCode } from '../../../errors/types/validation';
+import { assertValidationError } from '../../../errors/utils/assertValidationError';
+
+export const resolveIdentityId = (identityId?: string): string => {
+	assertValidationError(!!identityId, StorageValidationErrorCode.NoIdentityId);
+
+	return identityId;
+};