From ae40fec53a53b58c81d07b35d136f59a189e9a52 Mon Sep 17 00:00:00 2001
From: Manuel Iglesias <manuelig@amazon.com>
Date: Sat, 1 Dec 2018 14:23:28 -0800
Subject: [PATCH] feat(@aws-amplify/api,graphql): DocumentNode support for
 GraphQL Client

Fixes #1237
---
 packages/api/__tests__/API-test.ts | 69 ++++++++++++++++++++++++++++++
 packages/api/src/API.ts            | 15 +++----
 packages/api/src/types/index.ts    |  3 +-
 3 files changed, 77 insertions(+), 10 deletions(-)

diff --git a/packages/api/__tests__/API-test.ts b/packages/api/__tests__/API-test.ts
index 25a038c496c..22ea35021f4 100644
--- a/packages/api/__tests__/API-test.ts
+++ b/packages/api/__tests__/API-test.ts
@@ -107,6 +107,75 @@ describe('API test', () => {
             expect(spyon).toBeCalledWith(url, init);
         });
 
+        test('happy-case-query-ast', async () => {
+            const spyonAuth = jest.spyOn(Credentials, 'get').mockImplementationOnce(() => {
+                return new Promise((res, rej) => {
+                    res('cred');
+                });
+            });
+
+            const spyon = jest.spyOn(RestClient.prototype, 'post')
+                .mockImplementationOnce((url, init) => {
+                    return new Promise((res, rej) => {
+                        res({})
+                    });
+                });
+
+            const api = new API(config);
+            const url = 'https://appsync.amazonaws.com',
+                region = 'us-east-2',
+                apiKey = 'secret_api_key',
+                variables = { id: '809392da-ec91-4ef0-b219-5238a8f942b2' };
+            api.configure({
+                aws_appsync_graphqlEndpoint: url,
+                aws_appsync_region: region,
+                aws_appsync_authenticationType: 'API_KEY',
+                aws_appsync_apiKey: apiKey
+            })
+            const GetEvent = `query GetEvent($id: ID! $nextToken: String) {
+                getEvent(id: $id) {
+                    id
+                    name
+                    where
+                    when
+                    description
+                    comments(nextToken: $nextToken) {
+                      items {
+                        commentId
+                        content
+                        createdAt
+                      }
+                    }
+                  }
+                }`;
+
+            const doc = parse(GetEvent);
+            const query = print(doc);
+
+            const headers = {
+                Authorization: null,
+                'X-Api-Key': apiKey
+            };
+
+            const body = {
+                query,
+                variables,
+            };
+
+            const init = {
+                headers,
+                body,
+                signerServiceInfo: {
+                    service: 'appsync',
+                    region,
+                }
+            };
+
+            await api.graphql(graphqlOperation(doc, variables));
+
+            expect(spyon).toBeCalledWith(url, init);
+        });
+
         test('happy-case-query-oidc', async () => {
             const spyonAuth = jest.spyOn(Credentials, 'get').mockImplementationOnce(() => {
                 return new Promise((res, rej) => {
diff --git a/packages/api/src/API.ts b/packages/api/src/API.ts
index f85594c2eba..2287edbf2a6 100644
--- a/packages/api/src/API.ts
+++ b/packages/api/src/API.ts
@@ -325,11 +325,11 @@ export default class APIClass {
      * @param {GraphQLOptions} GraphQL Options
      * @returns {Promise<GraphQLResult> | Observable<object>}
      */
-    graphql({ query, variables = {} }: GraphQLOptions) {
+    graphql({ query: paramQuery, variables = {} }: GraphQLOptions) {
 
-        const doc = parse(query);
+        const query = typeof paramQuery === 'string' ? parse(paramQuery) : paramQuery;
 
-        const [operationDef = {},] = doc.definitions.filter(def => def.kind === 'OperationDefinition');
+        const [operationDef = {},] = query.definitions.filter(def => def.kind === 'OperationDefinition');
         const { operation: operationType } = operationDef as OperationDefinitionNode;
 
         switch (operationType) {
@@ -343,7 +343,7 @@ export default class APIClass {
         throw new Error(`invalid operation type: ${operationType}`);
     }
 
-    private async _graphql({ query: queryStr, variables }: GraphQLOptions, additionalHeaders = {})
+    private async _graphql({ query, variables }: GraphQLOptions, additionalHeaders = {})
         : Promise<GraphQLResult> {
         if (!this._api) {
             await this.createInstance();
@@ -357,20 +357,17 @@ export default class APIClass {
             graphql_endpoint_iam_region: customEndpointRegion
         } = this._options;
 
-        const doc = parse(queryStr);
-        const query = print(doc);
-
         const headers = {
             ...(!customGraphqlEndpoint && await this._headerBasedAuth()),
             ...(customGraphqlEndpoint &&
                 (customEndpointRegion ? await this._headerBasedAuth('AWS_IAM') : { Authorization: null })
             ),
             ...additionalHeaders,
-            ... await graphql_headers({ query: doc, variables }),
+            ... await graphql_headers({ query, variables }),
         };
 
         const body = {
-            query,
+            query: print(query),
             variables,
         };
 
diff --git a/packages/api/src/types/index.ts b/packages/api/src/types/index.ts
index 677d38b2ee7..f8c72be0efb 100644
--- a/packages/api/src/types/index.ts
+++ b/packages/api/src/types/index.ts
@@ -10,6 +10,7 @@
  * CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
+import { DocumentNode } from 'graphql/language/ast';
 
 /**
 * RestClient instance options
@@ -66,7 +67,7 @@ export interface apiOptions {
 }
 
 export interface GraphQLOptions {
-    query: string,
+    query: string | DocumentNode,
     variables?: object,
 }