fix: preserve ssr context when using DataStore (aws-amplify#10088)

.circleci/config.yml

@@ -1139,6 +1139,24 @@ jobs:
           spec: search-outside-map
           # Temp fix:
           browser: chrome
+
+  integ_next_datastore_owner_auth:
+    parameters:
+      browser:
+        type: string
+    executor: js-test-executor
+    <<: *test_env_vars
+    working_directory: ~/amplify-js-samples-staging/samples/next/datastore/owner-based-default
+    steps:
+      - prepare_test_env
+      - integ_test_js:
+          test_name: 'next owner auth'
+          framework: next
+          category: datastore
+          sample_name: owner-based-default
+          spec: next-owner-based-default
+          browser: << parameters.browser >>
+
   deploy:
     executor: macos-executor
     working_directory: ~/amplify-js
@@ -1583,6 +1601,15 @@ workflows:
           matrix:
             parameters:
               <<: *test_browsers
+      - integ_next_datastore_owner_auth:
+          requires:
+            - integ_setup
+            - build
+          filters:
+            <<: *releasable_branches
+          matrix:
+            parameters:
+              <<: *test_browsers
       - deploy:
           filters:
             <<: *releasable_branches
@@ -1630,6 +1657,7 @@ workflows:
             - integ_auth_test_cypress_no_ui
             - integ_react_graphql_api
             - integ_react_geo
+            - integ_next_datastore_owner_auth
       - post_release:
           filters:
             branches:

packages/aws-amplify/__tests__/withSSRContext-test.ts

@@ -1,4 +1,4 @@
-import { Amplify, CredentialsClass, UniversalStorage } from '@aws-amplify/core';
+import { Amplify, UniversalStorage } from '@aws-amplify/core';
 
 import { withSSRContext } from '../src/withSSRContext';
 
@@ -69,6 +69,16 @@ describe('withSSRContext', () => {
 		it('should be a different instance than Amplify.DataStore', () => {
 			expect(withSSRContext().DataStore).not.toBe(Amplify.DataStore);
 		});
+
+		it('should use Amplify components from the ssr context', () => {
+			const { Auth, API, DataStore } = withSSRContext();
+
+			expect(DataStore.Auth).toBe(Auth);
+			expect(DataStore.Auth).not.toBe(Amplify.Auth);
+
+			expect(DataStore.API).toBe(API);
+			expect(DataStore.API).not.toBe(Amplify.API);
+		});
 	});
 
 	describe('I18n', () => {

packages/datastore/__tests__/authStrategies.test.ts

@@ -449,9 +449,11 @@ async function testMultiAuthStrategy({
 }) {
 	mockCurrentUser({ hasAuthenticatedUser });
 
-	const multiAuthStrategy =
+	const multiAuthStrategyWrapper =
 		require('../src/authModeStrategies/multiAuthStrategy').multiAuthStrategy;
 
+	const multiAuthStrategy = multiAuthStrategyWrapper({});
+
 	const schema = getAuthSchema(authRules);
 
 	const authModes = await multiAuthStrategy({

packages/datastore/__tests__/mutation.test.ts

@@ -1,5 +1,4 @@
 const mockRestPost = jest.fn();
-import { RestClient } from '@aws-amplify/api-rest';
 import {
 	MutationProcessor,
 	safeJitteredBackoff,
@@ -371,7 +370,9 @@ async function instantiateMutationProcessor({
 			aws_appsync_apiKey: 'da2-xxxxxxxxxxxxxxxxxxxxxx',
 		},
 		() => null,
-		errorHandler
+		errorHandler,
+		() => null as any,
+		{} as any
 	);
 
 	(mutationProcessor as any).observer = true;

packages/datastore/__tests__/sync.test.ts

@@ -483,7 +483,8 @@ function jitteredRetrySyncProcessorSetup({
 		null, // syncPredicates
 		{ aws_appsync_authenticationType: 'userPools' },
 		defaultAuthStrategy,
-		errorHandler
+		errorHandler,
+		{}
 	);
 
 	return SyncProcessor;

packages/datastore/src/authModeStrategies/multiAuthStrategy.ts

@@ -5,6 +5,7 @@ import {
 	ModelAttributeAuthProperty,
 	ModelAttributeAuthProvider,
 	ModelAttributeAuthAllow,
+	AmplifyContext,
 } from '../types';
 
 function getProviderFromRule(
@@ -121,29 +122,31 @@ function getAuthRules({
 	return Array.from(authModes);
 }
 
-export const multiAuthStrategy: AuthModeStrategy = async ({
-	schema,
-	modelName,
-}) => {
-	let currentUser;
-	try {
-		currentUser = await Auth.currentAuthenticatedUser();
-	} catch (e) {
-		// No current user
-	}
+export const multiAuthStrategy: (
+	amplifyContext: AmplifyContext
+) => AuthModeStrategy =
+	(amplifyContext: AmplifyContext) =>
+	async ({ schema, modelName }) => {
+		amplifyContext.Auth = amplifyContext.Auth || Auth;
+		let currentUser;
+		try {
+			currentUser = await amplifyContext.Auth.currentAuthenticatedUser();
+		} catch (e) {
+			// No current user
+		}
 
-	const { attributes } = schema.namespaces.user.models[modelName];
+		const { attributes } = schema.namespaces.user.models[modelName];
 
-	if (attributes) {
-		const authAttribute = attributes.find(attr => attr.type === 'auth');
+		if (attributes) {
+			const authAttribute = attributes.find(attr => attr.type === 'auth');
 
-		if (authAttribute.properties && authAttribute.properties.rules) {
-			const sortedRules = sortAuthRulesWithPriority(
-				authAttribute.properties.rules
-			);
+			if (authAttribute?.properties?.rules) {
+				const sortedRules = sortAuthRulesWithPriority(
+					authAttribute.properties.rules
+				);
 
-			return getAuthRules({ currentUser, rules: sortedRules });
+				return getAuthRules({ currentUser, rules: sortedRules });
+			}
 		}
-	}
-	return [];
-};
+		return [];
+	};

packages/datastore/src/datastore/datastore.ts

@@ -1,4 +1,7 @@
+import API from '@aws-amplify/api';
 import { Amplify, ConsoleLogger as Logger, Hub, JS } from '@aws-amplify/core';
+import { Auth } from '@aws-amplify/auth';
+import Cache from '@aws-amplify/cache';
 import {
 	Draft,
 	immerable,
@@ -54,6 +57,7 @@ import {
 	isNonModelFieldType,
 	isModelFieldType,
 	ObserveQueryOptions,
+	AmplifyContext,
 } from '../types';
 import {
 	DATASTORE,
@@ -710,6 +714,11 @@ function getNamespace(): SchemaNamespace {
 }
 
 class DataStore {
+	// reference to configured category instances. Used for preserving SSR context
+	Auth = Auth;
+	API = API;
+	Cache = Cache;
+
 	private amplifyConfig: Record<string, any> = {};
 	private authModeStrategy: AuthModeStrategy;
 	private conflictHandler: ConflictHandler;
@@ -727,6 +736,12 @@ class DataStore {
 		new WeakMap<SchemaModel, ModelPredicate<any>>();
 	private sessionId: string;
 	private storageAdapter: Adapter;
+	// object that gets passed to descendent classes. Allows us to pass these down by reference
+	private amplifyContext: AmplifyContext = {
+		Auth: this.Auth,
+		API: this.API,
+		Cache: this.Cache,
+	};
 
 	getModuleName() {
 		return 'DataStore';
@@ -777,7 +792,8 @@ class DataStore {
 				this.errorHandler,
 				this.syncPredicates,
 				this.amplifyConfig,
-				this.authModeStrategy
+				this.authModeStrategy,
+				this.amplifyContext
 			);
 
 			// tslint:disable-next-line:max-line-length
@@ -1417,6 +1433,10 @@ class DataStore {
 	};
 
 	configure = (config: DataStoreConfig = {}) => {
+		this.amplifyContext.Auth = this.Auth;
+		this.amplifyContext.API = this.API;
+		this.amplifyContext.Cache = this.Cache;
+
 		const {
 			DataStore: configDataStore,
 			authModeStrategyType: configAuthModeStrategyType,
@@ -1446,7 +1466,7 @@ class DataStore {
 
 		switch (authModeStrategyType) {
 			case AuthModeStrategyType.MULTI_AUTH:
-				this.authModeStrategy = multiAuthStrategy;
+				this.authModeStrategy = multiAuthStrategy(this.amplifyContext);
 				break;
 			case AuthModeStrategyType.DEFAULT:
 				this.authModeStrategy = defaultAuthStrategy;

packages/datastore/src/sync/index.ts

@@ -21,6 +21,7 @@ import {
 	TypeConstructorMap,
 	ModelPredicate,
 	AuthModeStrategy,
+	AmplifyContext,
 } from '../types';
 import { exhaustiveCheck, getNow, SYNC, USER } from '../util';
 import DataStoreConnectivity from './datastoreConnectivity';
@@ -117,7 +118,8 @@ export class SyncEngine {
 		errorHandler: ErrorHandler,
 		private readonly syncPredicates: WeakMap<SchemaModel, ModelPredicate<any>>,
 		private readonly amplifyConfig: Record<string, any> = {},
-		private readonly authModeStrategy: AuthModeStrategy
+		private readonly authModeStrategy: AuthModeStrategy,
+		private readonly amplifyContext: AmplifyContext
 	) {
 		const MutationEvent = this.modelClasses[
 			'MutationEvent'
@@ -137,15 +139,19 @@ export class SyncEngine {
 			this.syncPredicates,
 			this.amplifyConfig,
 			this.authModeStrategy,
-			errorHandler
+			errorHandler,
+			this.amplifyContext
 		);
+
 		this.subscriptionsProcessor = new SubscriptionProcessor(
 			this.schema,
 			this.syncPredicates,
 			this.amplifyConfig,
 			this.authModeStrategy,
-			errorHandler
+			errorHandler,
+			this.amplifyContext
 		);
+
 		this.mutationsProcessor = new MutationProcessor(
 			this.schema,
 			this.storage,
@@ -156,8 +162,10 @@ export class SyncEngine {
 			this.amplifyConfig,
 			this.authModeStrategy,
 			errorHandler,
-			conflictHandler
+			conflictHandler,
+			this.amplifyContext
 		);
+
 		this.datastoreConnectivity = new DataStoreConnectivity();
 	}
 

packages/datastore/src/sync/processors/mutation.ts

@@ -25,6 +25,7 @@ import {
 	SchemaModel,
 	TypeConstructorMap,
 	ProcessName,
+	AmplifyContext,
 } from '../../types';
 import { exhaustiveCheck, USER, USER_AGENT_SUFFIX_DATASTORE } from '../../util';
 import { MutationEventOutbox } from '../outbox';
@@ -66,8 +67,10 @@ class MutationProcessor {
 		private readonly amplifyConfig: Record<string, any> = {},
 		private readonly authModeStrategy: AuthModeStrategy,
 		private readonly errorHandler: ErrorHandler,
-		private readonly conflictHandler?: ConflictHandler
+		private readonly conflictHandler: ConflictHandler,
+		private readonly amplifyContext: AmplifyContext
 	) {
+		this.amplifyContext.API = this.amplifyContext.API || API;
 		this.generateQueries();
 	}
 
@@ -282,7 +285,7 @@ class MutationProcessor {
 				do {
 					try {
 						const result = <GraphQLResult<Record<string, PersistentModel>>>(
-							await API.graphql(tryWith)
+							await this.amplifyContext.API.graphql(tryWith)
 						);
 						return [result, opName, modelDefinition];
 					} catch (err) {
@@ -349,7 +352,7 @@ class MutationProcessor {
 
 									const serverData = <
 										GraphQLResult<Record<string, PersistentModel>>
-									>await API.graphql({
+									>await this.amplifyContext.API.graphql({
 										query,
 										variables: { id: variables.input.id },
 										authMode,