Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(api): add custom endpoint support to API #14086

Merged
merged 17 commits into from
Dec 23, 2024
Merged
Show file tree
Hide file tree
Changes from 11 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 23 additions & 25 deletions packages/adapter-nextjs/src/api/generateServerClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,10 @@ import {
V6ClientSSRRequest,
} from '@aws-amplify/api-graphql';
import {
GraphQLAuthMode,
parseAmplifyConfig,
} from '@aws-amplify/core/internals/utils';
CommonPublicClientOptions,
DefaultCommonClientOptions,
} from '@aws-amplify/api-graphql/dist/esm/internals/types';
svidgen marked this conversation as resolved.
Show resolved Hide resolved
import { parseAmplifyConfig } from '@aws-amplify/core/internals/utils';

import { NextServer } from '../types';

Expand All @@ -23,14 +24,10 @@ import { createServerRunnerForAPI } from './createServerRunnerForAPI';
interface CookiesClientParams {
cookies: NextServer.ServerComponentContext['cookies'];
config: NextServer.CreateServerRunnerInput['config'];
authMode?: GraphQLAuthMode;
authToken?: string;
}

interface ReqClientParams {
config: NextServer.CreateServerRunnerInput['config'];
authMode?: GraphQLAuthMode;
authToken?: string;
}

/**
Expand All @@ -44,13 +41,10 @@ interface ReqClientParams {
*/
export function generateServerClientUsingCookies<
T extends Record<any, any> = never,
>({
config,
cookies,
authMode,
authToken,
}: CookiesClientParams): V6ClientSSRCookies<T> {
if (typeof cookies !== 'function') {
Options extends CommonPublicClientOptions &
CookiesClientParams = DefaultCommonClientOptions & CookiesClientParams,
>(options: Options): V6ClientSSRCookies<T, Options> {
if (typeof options.cookies !== 'function') {
throw new AmplifyServerContextError({
message:
'generateServerClientUsingCookies is only compatible with the `cookies` Dynamic Function available in Server Components.',
Expand All @@ -61,24 +55,25 @@ export function generateServerClientUsingCookies<
}

const { runWithAmplifyServerContext, resourcesConfig } =
createServerRunnerForAPI({ config });
createServerRunnerForAPI({ config: options.config });

// This function reference gets passed down to InternalGraphQLAPI.ts.graphql
// where this._graphql is passed in as the `fn` argument
// causing it to always get invoked inside `runWithAmplifyServerContext`
const getAmplify = (fn: (amplify: any) => Promise<any>) =>
runWithAmplifyServerContext({
nextServerContext: { cookies },
nextServerContext: { cookies: options.cookies },
operation: contextSpec =>
fn(getAmplifyServerContext(contextSpec).amplify),
});

return generateClientWithAmplifyInstance<T, V6ClientSSRCookies<T>>({
const { cookies: _cookies, config: _config, ...params } = options;

return generateClientWithAmplifyInstance<T, V6ClientSSRCookies<T, Options>>({
amplify: getAmplify,
config: resourcesConfig,
authMode,
authToken,
});
...params,
} as any); // TS can't narrow the type here.
}

/**
Expand All @@ -99,12 +94,15 @@ export function generateServerClientUsingCookies<
*/
export function generateServerClientUsingReqRes<
T extends Record<any, any> = never,
>({ config, authMode, authToken }: ReqClientParams): V6ClientSSRRequest<T> {
const amplifyConfig = parseAmplifyConfig(config);
Options extends CommonPublicClientOptions &
ReqClientParams = DefaultCommonClientOptions & ReqClientParams,
>(options: Options): V6ClientSSRRequest<T, Options> {
const amplifyConfig = parseAmplifyConfig(options.config);

const { config: _config, ...params } = options;

return generateClient<T>({
config: amplifyConfig,
authMode,
authToken,
});
...params,
}) as any;
}
54 changes: 48 additions & 6 deletions packages/api-graphql/__tests__/internals/generateClient.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -270,7 +270,14 @@ describe('generateClient', () => {
const client = generateClient<Schema>({ amplify: Amplify });

const spy = jest.fn(() => from([graphqlMessage]));
(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
(raw.GraphQLAPI as any).appSyncRealTime = {
get() {
return { subscribe: spy }
},
set() {
// not needed for test mock
}
};

expect(normalizePostGraphqlCalls(spy)).toMatchSnapshot();

Expand Down Expand Up @@ -497,7 +504,14 @@ describe('generateClient', () => {
const client = generateClient<Schema>({ amplify: Amplify });

const spy = jest.fn(() => from([graphqlMessage]));
(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
(raw.GraphQLAPI as any).appSyncRealTime = {
get() {
return { subscribe: spy }
},
set() {
// not needed for test mock
}
};

client.models.Note.onCreate({
filter: graphqlVariables.filter,
Expand Down Expand Up @@ -531,7 +545,14 @@ describe('generateClient', () => {
});

const spy = jest.fn(() => from([graphqlMessage]));
(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
(raw.GraphQLAPI as any).appSyncRealTime = {
get() {
return { subscribe: spy }
},
set() {
// not needed for test mock
}
};

client.models.Note.onCreate({
filter: graphqlVariables.filter,
Expand Down Expand Up @@ -561,7 +582,14 @@ describe('generateClient', () => {
const client = generateClient<Schema>({ amplify: Amplify });

const spy = jest.fn(() => from([graphqlMessage]));
(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
(raw.GraphQLAPI as any).appSyncRealTime = {
get() {
return { subscribe: spy }
},
set() {
// not needed for test mock
}
};

client.models.Note.onCreate({
filter: graphqlVariables.filter,
Expand All @@ -583,7 +611,14 @@ describe('generateClient', () => {
const client = generateClient<Schema>({ amplify: Amplify });

const spy = jest.fn(() => from([graphqlMessage]));
(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
(raw.GraphQLAPI as any).appSyncRealTime = {
get() {
return { subscribe: spy }
},
set() {
// not needed for test mock
}
};

client.models.Note.onCreate({
filter: graphqlVariables.filter,
Expand Down Expand Up @@ -711,7 +746,14 @@ describe('generateClient', () => {
const client = generateClient<Schema>({ amplify: Amplify });

const spy = jest.fn(() => from([graphqlMessage]));
(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
(raw.GraphQLAPI as any).appSyncRealTime = {
get() {
return { subscribe: spy }
},
set() {
// not needed for test mock
}
};

client.models.Note.onCreate({
filter: graphqlVariables.filter,
Expand Down
61 changes: 46 additions & 15 deletions packages/api-graphql/src/internals/InternalGraphQLAPI.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ export class InternalGraphQLAPIClass {
/**
* @private
*/
private appSyncRealTime = new AWSAppSyncRealTimeProvider();
private appSyncRealTime = new Map<string, AWSAppSyncRealTimeProvider>();

private _api = {
post,
Expand Down Expand Up @@ -88,7 +88,14 @@ export class InternalGraphQLAPIClass {
amplify:
| AmplifyClassV6
| ((fn: (amplify: any) => Promise<any>) => Promise<AmplifyClassV6>),
{ query: paramQuery, variables = {}, authMode, authToken }: GraphQLOptions,
{
query: paramQuery,
variables = {},
authMode,
authToken,
endpoint,
apiKey,
}: GraphQLOptions,
additionalHeaders?: CustomHeaders,
customUserAgentDetails?: CustomUserAgentDetails,
): Observable<GraphQLResult<T>> | Promise<GraphQLResult<T>> {
Expand All @@ -115,7 +122,7 @@ export class InternalGraphQLAPIClass {
if (isAmplifyInstance(amplify)) {
responsePromise = this._graphql<T>(
amplify,
{ query, variables, authMode },
{ query, variables, authMode, apiKey, endpoint },
headers,
abortController,
customUserAgentDetails,
Expand All @@ -127,7 +134,7 @@ export class InternalGraphQLAPIClass {
const wrapper = async (amplifyInstance: AmplifyClassV6) => {
const result = await this._graphql<T>(
amplifyInstance,
{ query, variables, authMode },
{ query, variables, authMode, apiKey, endpoint },
headers,
abortController,
customUserAgentDetails,
Expand All @@ -152,7 +159,7 @@ export class InternalGraphQLAPIClass {
case 'subscription':
return this._graphqlSubscribe(
amplify as AmplifyClassV6,
{ query, variables, authMode },
{ query, variables, authMode, apiKey, endpoint },
headers,
customUserAgentDetails,
authToken,
Expand All @@ -164,7 +171,13 @@ export class InternalGraphQLAPIClass {

private async _graphql<T = any>(
amplify: AmplifyClassV6,
{ query, variables, authMode: explicitAuthMode }: GraphQLOptions,
{
query,
variables,
authMode: authModeOverride,
endpoint: endpointOverride,
apiKey: apiKeyOverride,
}: GraphQLOptions,
additionalHeaders: CustomHeaders = {},
abortController: AbortController,
customUserAgentDetails?: CustomUserAgentDetails,
Expand All @@ -179,7 +192,7 @@ export class InternalGraphQLAPIClass {
defaultAuthMode,
} = resolveConfig(amplify);

const initialAuthMode = explicitAuthMode || defaultAuthMode || 'iam';
const initialAuthMode = authModeOverride || defaultAuthMode || 'iam';
// identityPool is an alias for iam. TODO: remove 'iam' in v7
const authMode =
initialAuthMode === 'identityPool' ? 'iam' : initialAuthMode;
Expand All @@ -205,7 +218,7 @@ export class InternalGraphQLAPIClass {
const requestOptions: RequestOptions = {
method: 'POST',
url: new AmplifyUrl(
customEndpoint || appSyncGraphqlEndpoint || '',
endpointOverride || customEndpoint || appSyncGraphqlEndpoint || '',
).toString(),
queryString: print(query as DocumentNode),
};
Expand All @@ -226,7 +239,7 @@ export class InternalGraphQLAPIClass {
const authHeaders = await headerBasedAuth(
amplify,
authMode,
apiKey,
apiKeyOverride ?? apiKey,
additionalCustomHeaders,
);

Expand Down Expand Up @@ -282,7 +295,8 @@ export class InternalGraphQLAPIClass {
};
}

const endpoint = customEndpoint || appSyncGraphqlEndpoint;
const endpoint =
endpointOverride || customEndpoint || appSyncGraphqlEndpoint;

if (!endpoint) {
throw createGraphQLResultWithError<T>(new GraphQLApiError(NO_ENDPOINT));
Expand Down Expand Up @@ -341,15 +355,21 @@ export class InternalGraphQLAPIClass {

private _graphqlSubscribe(
amplify: AmplifyClassV6,
{ query, variables, authMode: explicitAuthMode }: GraphQLOptions,
{
query,
variables,
authMode: authModeOverride,
apiKey: apiKeyOverride,
endpoint,
}: GraphQLOptions,
additionalHeaders: CustomHeaders = {},
customUserAgentDetails?: CustomUserAgentDetails,
authToken?: string,
): Observable<any> {
const config = resolveConfig(amplify);

const initialAuthMode =
explicitAuthMode || config?.defaultAuthMode || 'iam';
authModeOverride || config?.defaultAuthMode || 'iam';
// identityPool is an alias for iam. TODO: remove 'iam' in v7
const authMode =
initialAuthMode === 'identityPool' ? 'iam' : initialAuthMode;
Expand All @@ -364,15 +384,26 @@ export class InternalGraphQLAPIClass {
*/
const { headers: libraryConfigHeaders } = resolveLibraryOptions(amplify);

return this.appSyncRealTime
const appSyncGraphqlEndpoint = endpoint ?? config?.endpoint;

// TODO: This could probably be an exception. But, lots of tests rely on
// attempting to connect to nowhere. So, I'm treating as the opposite of
// a Chesterton's fence for now. (A fence I shouldn't build, because I don't
// know why somethings depends on its absence!)
const memoKey = appSyncGraphqlEndpoint ?? 'none';
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this memoKey sufficient? Wondering if a websocket might ever be accessed with different authtypes to achieve different permissions where both would have the same endpoint (or something similar)

I think in all case, overriding a graphql calls authtype works, but maybe a client with authtype A will be recreated with authtype B and the dev will be surprised that they get the original instance with type A...

Any other config that might be needed in the memoKey to avoid surprises?

Copy link
Member Author

@svidgen svidgen Dec 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it should be OK. We have normally used a single socket for all connections regardless of the authmode needed by the subscription itself. In chatting with Ivan about how to do this, I think he leaned towards keeping connections pooled by client — in most cases, I think that's slightly better for customers. Fewer local resources consumed at the very least, etc.. It's marginal. I wouldn't die on that hill.

If it causes issues, I think this behavior is pretty OK for us to change later.

const realtimeProvider =
this.appSyncRealTime.get(memoKey) ?? new AWSAppSyncRealTimeProvider();
this.appSyncRealTime.set(memoKey, realtimeProvider);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This overwrites the memoKey, even when we just retrieved the provider from the set doesn't it? It's probably not doing harm, but seems like an unnecessary operation.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah. I weighed it against a conditional check to do the same. This seemed simpler to read. But, I don't mind changing it if you disagree.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The conditional check seems like fewer moving parts or things to go wrong, but I also don't think I can come up with an example of when this would fail us, so this is a nit at best.


return realtimeProvider
.subscribe(
{
query: print(query as DocumentNode),
variables,
appSyncGraphqlEndpoint: config?.endpoint,
appSyncGraphqlEndpoint,
region: config?.region,
authenticationType: authMode,
apiKey: config?.apiKey,
apiKey: apiKeyOverride ?? config?.apiKey,
additionalHeaders,
authToken,
libraryConfigHeaders,
Expand Down
Loading
Loading