Added ability to retry unauthorized requests

Author: maneesht

packages/data-connect/CHANGELOG.md

@@ -0,0 +1,5 @@
+# @firebase/data-connect
+## UNRELEASED
+* Updated reporting to use @firebase/data-connect instead of @firebase/connect
+* Added functionality to retry queries and mutations if the server responds with UNAUTHENTICATED.
+

packages/data-connect/emulator.json

@@ -0,0 +1 @@
+{"service_id":"l","schema":{"files":[{"path":"schema/queries.schema.gql","content":"type Post @table {content: String!}"}]},"connectors":{"c":{"files":[{"path":"operations/post.gql","content":"query getPost($id: UUID!) @auth(level: PUBLIC) {\n  post(id: $id) {\n    content\n  }\n}\nquery listPosts @auth(level: PUBLIC) {\n  posts {\n    id,\n    content\n  }\n}\n\n"},{"path":"operations/mutations.gql","content":"mutation seedDatabase($id: UUID!, $content: String!) @auth(level: PUBLIC) {\n    post: post_insert(data: {id: $id, content: $content})\n}\nmutation removePost($id: UUID!) @auth(level: PUBLIC) {\n    post: post_delete(id: $id)\n}"}]}},"connection_string":"postgresql://postgres:secretpassword@localhost:5432/postgres?sslmode=disable"}

packages/data-connect/package.json

@@ -38,6 +38,7 @@
     "test:all": "npm run test:node",
     "test:browser": "karma start --single-run",
     "test:node": "TS_NODE_FILES=true TS_NODE_CACHE=NO TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' nyc --reporter lcovonly -- mocha 'test/{,!(browser)/**/}*.test.ts' --file src/index.node.ts --config ../../config/mocharc.node.js",
+    "test:unit": "TS_NODE_FILES=true TS_NODE_CACHE=NO TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' nyc --reporter lcovonly -- mocha 'test/unit/*.test.ts' --file src/index.node.ts --config ../../config/mocharc.node.js",
     "test:emulator": "ts-node --compiler-options='{\"module\":\"commonjs\"}' ../../scripts/emulator-testing/dataconnect-test-runner.ts",
     "api-report": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' ts-node ../../repo-scripts/prune-dts/extract-public-api.ts --package data-connect --packageRoot . --typescriptDts ./dist/src/index.d.ts --rollupDts ./dist/private.d.ts --untrimmedRollupDts ./dist/internal.d.ts --publicDts ./dist/public.d.ts && yarn api-report:api-json",
     "api-report:api-json": "rm -rf temp && api-extractor run --local --verbose",

packages/data-connect/src/core/error.ts

@@ -23,7 +23,8 @@ export type DataConnectErrorCode =
   | 'not-initialized'
   | 'not-supported'
   | 'invalid-argument'
-  | 'partial-error';
+  | 'partial-error'
+  | 'unauthorized';
 
 export type Code = DataConnectErrorCode;
 
@@ -33,7 +34,8 @@ export const Code = {
   NOT_INITIALIZED: 'not-initialized' as DataConnectErrorCode,
   NOT_SUPPORTED: 'not-supported' as DataConnectErrorCode,
   INVALID_ARGUMENT: 'invalid-argument' as DataConnectErrorCode,
-  PARTIAL_ERROR: 'partial-error' as DataConnectErrorCode
+  PARTIAL_ERROR: 'partial-error' as DataConnectErrorCode,
+  UNAUTHORIZED: 'unauthorized' as DataConnectErrorCode
 };
 
 /** An error returned by a DataConnect operation. */

packages/data-connect/src/network/fetch.ts

@@ -44,6 +44,7 @@ export function dcFetch<T, U>(
   }
   const bodyStr = JSON.stringify(body);
   logDebug(`Making request out to ${url} with body: ${bodyStr}`);
+  
   return connectFetch(url, {
     body: bodyStr,
     method: 'POST',
@@ -64,9 +65,12 @@ export function dcFetch<T, U>(
         throw new DataConnectError(Code.OTHER, JSON.stringify(e));
       }
       if (response.status >= 400) {
-        logError(
-          'Error while performing request: ' + JSON.stringify(jsonResponse)
-        );
+        // logError(
+        //   'Error while performing request: ' + JSON.stringify(jsonResponse)
+        // );
+        if(response.status === 401) {
+          throw new DataConnectError(Code.UNAUTHORIZED, JSON.stringify(jsonResponse));
+        }
         throw new DataConnectError(Code.OTHER, JSON.stringify(jsonResponse));
       }
       return jsonResponse;

packages/data-connect/src/network/transport/rest.ts

@@ -20,7 +20,7 @@ import { DataConnectError, Code } from '../../core/error';
 import { AuthTokenProvider } from '../../core/FirebaseAuthProvider';
 import { logDebug } from '../../logger';
 import { addToken, urlBuilder } from '../../util/url';
-import { dcFetch } from '../fetch';
+import { dcFetch, initializeFetch } from '../fetch';
 
 import { DataConnectTransport } from '.';
 
@@ -34,6 +34,7 @@ export class RESTTransport implements DataConnectTransport {
   private _serviceName: string;
   private _accessToken: string | null = null;
   private _authInitialized = false;
+  private _lastToken: string | null = null;
   constructor(
     options: DataConnectOptions,
     private apiKey?: string | undefined,
@@ -93,14 +94,14 @@ export class RESTTransport implements DataConnectTransport {
     this._accessToken = newToken;
   }
 
-  getWithAuth() {
+  getWithAuth(forceToken = false) {
     let starterPromise: Promise<string | null> = new Promise(resolve =>
       resolve(this._accessToken)
     );
     if (!this._authInitialized) {
       if (this.authProvider) {
         starterPromise = this.authProvider
-          .getToken(/*forceToken=*/ false)
+          .getToken(/*forceToken=*/ forceToken)
           .then(data => {
             if (!data) {
               return null;
@@ -115,13 +116,39 @@ export class RESTTransport implements DataConnectTransport {
     return starterPromise;
   }
 
+  _setLastToken(lastToken: string | null) {
+    this._lastToken = lastToken;
+  }
+
+  //TODO(mtewani): Remove any
+  withRetry<T>(promiseFactory: () => Promise<{ data: T, errors: Error[]}>, retry = false) {
+    let isNewToken = false;
+    return this.getWithAuth(retry)
+      .then(res => {
+        isNewToken = this._lastToken !== res;
+        this._lastToken = res;
+        return res;
+      })
+      .then(promiseFactory)
+      .catch(err => {
+        // Only retry if the result is unauthorized and the last token isn't the same as the new one.
+        if (
+          'code' in err &&
+          err.code === Code.UNAUTHORIZED &&
+          !retry && isNewToken
+        ) {
+          logDebug('Retrying due to unauthorized');
+          return this.withRetry(promiseFactory, true);
+        }
+        throw err;
+      });
+  }
+
   // TODO(mtewani): Update U to include shape of body defined in line 13.
   invokeQuery = <T, U = unknown>(queryName: string, body: U) => {
     const abortController = new AbortController();
-
     // TODO(mtewani): Update to proper value
-    const withAuth = this.getWithAuth().then(() => {
-      return dcFetch<T, U>(
+    const withAuth = this.withRetry(() => dcFetch<T, U>(
         addToken(`${this.endpointUrl}:executeQuery`, this.apiKey),
         {
           name: `projects/${this._project}/locations/${this._location}/services/${this._serviceName}/connectors/${this._connectorName}`,
@@ -130,16 +157,15 @@ export class RESTTransport implements DataConnectTransport {
         } as unknown as U, // TODO(mtewani): This is a patch, fix this.
         abortController,
         this._accessToken
-      );
-    });
+      ));
 
     return {
       then: withAuth.then.bind(withAuth)
     };
   };
   invokeMutation = <T, U = unknown>(mutationName: string, body: U) => {
     const abortController = new AbortController();
-    const taskResult = this.getWithAuth().then(() => {
+    const taskResult = this.withRetry(() => {
       return dcFetch<T, U>(
         addToken(`${this.endpointUrl}:executeMutation`, this.apiKey),
         {

packages/data-connect/test/unit/queries.test.ts

@@ -0,0 +1,79 @@
+import { FirebaseAuthTokenData } from '@firebase/auth-interop-types';
+import {
+  AuthTokenListener,
+  AuthTokenProvider,
+  DataConnectOptions,
+  FirebaseAuthProvider
+} from '../../src';
+import { RESTTransport } from '../../src/network/transport/rest';
+import { initializeFetch } from '../../src/network/fetch';
+import { expect } from 'chai';
+import * as chai from 'chai';
+import chaiAsPromised from 'chai-as-promised';
+import * as sinon from 'sinon';
+chai.use(chaiAsPromised);
+const options: DataConnectOptions = {
+  connector: 'c',
+  location: 'l',
+  projectId: 'p',
+  service: 's'
+};
+const INITIAL_TOKEN = 'initial token';
+class FakeAuthProvider implements AuthTokenProvider {
+  private token: string | null = INITIAL_TOKEN;
+  addTokenChangeListener(listener: AuthTokenListener): void {}
+  getToken(forceRefresh: boolean): Promise<FirebaseAuthTokenData | null> {
+    if (!forceRefresh) {
+      return Promise.resolve({ accessToken: this.token! });
+    }
+    return Promise.resolve({ accessToken: 'testToken' });
+  }
+  setToken(_token: string | null) {
+    this.token = _token;
+  }
+}
+const json = {
+  message: 'unauthorized'
+};
+
+const fakeFetchImpl = sinon.stub().returns(
+  Promise.resolve({
+    json: () => {
+      return Promise.resolve(json);
+    },
+    status: 401
+  } as Response)
+);
+describe('Queries', () => {
+  afterEach(() => {
+    fakeFetchImpl.resetHistory();
+  });
+  it('[QUERY] should retry auth whenever the fetcher returns with unauthorized', async () => {
+    initializeFetch(fakeFetchImpl);
+    const authProvider = new FakeAuthProvider();
+    const rt = new RESTTransport(options, undefined, authProvider);
+    await expect(
+      rt.invokeQuery('test', null)
+    ).to.eventually.be.rejectedWith(JSON.stringify(json));
+    expect(fakeFetchImpl.callCount).to.eq(2);
+  });
+  it('[MUTATION] should retry auth whenever the fetcher returns with unauthorized', async () => {
+    initializeFetch(fakeFetchImpl);
+    const authProvider = new FakeAuthProvider();
+    const rt = new RESTTransport(options, undefined, authProvider);
+    await expect(
+      rt.invokeMutation('test', null)
+    ).to.eventually.be.rejectedWith(JSON.stringify(json));
+    expect(fakeFetchImpl.callCount).to.eq(2);
+  });
+  it("should not retry auth whenever the fetcher returns with unauthorized and the token doesn't change", async () => {
+    initializeFetch(fakeFetchImpl);
+    const authProvider = new FakeAuthProvider();
+    const rt = new RESTTransport(options, undefined, authProvider);
+    rt._setLastToken('initial token');
+    await expect(
+      rt.invokeQuery('test', null) as Promise<any>
+    ).to.eventually.be.rejectedWith(JSON.stringify(json));
+    expect(fakeFetchImpl.callCount).to.eq(1);
+  });
+});