Add forceWebSockets() and forceLongPolling() (#6171)

* Implemented forceWebSockets() and forceLongPolling()

Author: maneesht

.changeset/lucky-items-wave.md

@@ -0,0 +1,6 @@
+---
+"@firebase/database-compat": minor
+"@firebase/database": minor
+---
+
+Add `forceWebSockets()` and `forceLongPolling()`

common/api-review/database.api.md

@@ -64,6 +64,12 @@ export function equalTo(value: number | string | boolean | null, key?: string):
 // @public
 export type EventType = 'value' | 'child_added' | 'child_changed' | 'child_moved' | 'child_removed';
 
+// @public
+export function forceLongPolling(): void;
+
+// @public
+export function forceWebSockets(): void;
+
 // @public
 export function get(query: Query): Promise<DataSnapshot>;
 

packages/database-compat/src/api/Database.ts

@@ -19,6 +19,8 @@
 import { FirebaseApp } from '@firebase/app-types';
 import { FirebaseService } from '@firebase/app-types/private';
 import {
+  forceLongPolling,
+  forceWebSockets,
   goOnline,
   connectDatabaseEmulator,
   goOffline,
@@ -51,7 +53,9 @@ export class Database implements FirebaseService, Compat<ModularDatabase> {
   constructor(readonly _delegate: ModularDatabase, readonly app: FirebaseApp) {}
 
   INTERNAL = {
-    delete: () => this._delegate._delete()
+    delete: () => this._delegate._delete(),
+    forceWebSockets,
+    forceLongPolling
   };
 
   /**

packages/database/src/api.standalone.ts

@@ -22,6 +22,8 @@ export {
   enableLogging,
   goOffline,
   goOnline,
+  forceWebSockets,
+  forceLongPolling,
   connectDatabaseEmulator
 } from './api/Database';
 export {

packages/database/src/api/Database.ts

@@ -41,11 +41,15 @@ import { RepoInfo } from '../core/RepoInfo';
 import { parseRepoInfo } from '../core/util/libs/parser';
 import { newEmptyPath, pathIsEmpty } from '../core/util/Path';
 import {
+  warn,
   fatal,
   log,
   enableLogging as enableLoggingImpl
 } from '../core/util/util';
 import { validateUrl } from '../core/util/validation';
+import { BrowserPollConnection } from '../realtime/BrowserPollConnection';
+import { TransportManager } from '../realtime/TransportManager';
+import { WebSocketConnection } from '../realtime/WebSocketConnection';
 
 import { ReferenceImpl } from './Reference_impl';
 
@@ -271,6 +275,31 @@ export class Database implements _FirebaseService {
   }
 }
 
+function checkTransportInit() {
+  if (TransportManager.IS_TRANSPORT_INITIALIZED) {
+    warn(
+      'Transport has already been initialized. Please call this function before calling ref or setting up a listener'
+    );
+  }
+}
+
+/**
+ * Force the use of websockets instead of longPolling.
+ */
+export function forceWebSockets() {
+  checkTransportInit();
+  BrowserPollConnection.forceDisallow();
+}
+
+/**
+ * Force the use of longPolling instead of websockets. This will be ignored if websocket protocol is used in databaseURL.
+ */
+export function forceLongPolling() {
+  checkTransportInit();
+  WebSocketConnection.forceDisallow();
+  BrowserPollConnection.forceAllow();
+}
+
 /**
  * Returns the instance of the Realtime Database SDK that is associated
  * with the provided {@link @firebase/app#FirebaseApp}. Initializes a new instance with

packages/database/src/realtime/BrowserPollConnection.ts

@@ -248,7 +248,7 @@ export class BrowserPollConnection implements Transport {
     this.addDisconnectPingFrame(this.id, this.password);
   }
 
-  private static forceAllow_: boolean;
+  static forceAllow_: boolean;
 
   /**
    * Forces long polling to be considered as a potential transport
@@ -257,7 +257,7 @@ export class BrowserPollConnection implements Transport {
     BrowserPollConnection.forceAllow_ = true;
   }
 
-  private static forceDisallow_: boolean;
+  static forceDisallow_: boolean;
 
   /**
    * Forces longpolling to not be considered as a potential transport

packages/database/src/realtime/TransportManager.ts

@@ -32,10 +32,21 @@ import { WebSocketConnection } from './WebSocketConnection';
 export class TransportManager {
   private transports_: TransportConstructor[];
 
+  // Keeps track of whether the TransportManager has already chosen a transport to use
+  static globalTransportInitialized_ = false;
+
   static get ALL_TRANSPORTS() {
     return [BrowserPollConnection, WebSocketConnection];
   }
 
+  /**
+   * Returns whether transport has been selected to ensure WebSocketConnection or BrowserPollConnection are not called after
+   * TransportManager has already set up transports_
+   */
+  static get IS_TRANSPORT_INITIALIZED() {
+    return this.globalTransportInitialized_;
+  }
+
   /**
    * @param repoInfo - Metadata around the namespace we're connecting to
    */
@@ -68,6 +79,7 @@ export class TransportManager {
           transports.push(transport);
         }
       }
+      TransportManager.globalTransportInitialized_ = true;
     }
   }
 

packages/database/test/transport.test.ts

@@ -0,0 +1,76 @@
+/**
+ * @license
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { CONSTANTS } from '@firebase/util';
+import { expect, use } from 'chai';
+import { createSandbox, SinonSandbox, SinonSpy } from 'sinon';
+import sinonChai from 'sinon-chai';
+
+import { forceLongPolling, forceWebSockets } from '../src';
+import * as Util from '../src/core/util/util';
+import { BrowserPollConnection } from '../src/realtime/BrowserPollConnection';
+import { TransportManager } from '../src/realtime/TransportManager';
+import { WebSocketConnection } from '../src/realtime/WebSocketConnection';
+
+use(sinonChai);
+const transportInitError =
+  'Transport has already been initialized. Please call this function before calling ref or setting up a listener';
+describe('Force Transport', () => {
+  const oldNodeValue = CONSTANTS.NODE_CLIENT;
+  let mySandbox: SinonSandbox;
+  let spyWarn: SinonSpy;
+  beforeEach(() => {
+    CONSTANTS.NODE_CLIENT = false;
+    mySandbox = createSandbox();
+    spyWarn = mySandbox.spy(Util, 'warn');
+  });
+  afterEach(() => {
+    // Resetting to old values
+    TransportManager.globalTransportInitialized_ = false;
+    CONSTANTS.NODE_CLIENT = oldNodeValue;
+    BrowserPollConnection.forceAllow_ = false;
+    BrowserPollConnection.forceDisallow_ = true;
+    WebSocketConnection.forceDisallow_ = false;
+    mySandbox.restore();
+  });
+  it('should enable websockets and disable longPolling', () => {
+    forceWebSockets();
+    expect(spyWarn.called).to.equal(false);
+    expect(WebSocketConnection.isAvailable()).to.equal(true);
+    expect(BrowserPollConnection.isAvailable()).to.equal(false);
+  });
+  it('should throw an error when calling forceWebsockets() if TransportManager has already been initialized', () => {
+    TransportManager.globalTransportInitialized_ = true;
+    forceWebSockets();
+    expect(spyWarn).to.have.been.calledWith(transportInitError);
+    expect(WebSocketConnection.isAvailable()).to.equal(true);
+    expect(BrowserPollConnection.isAvailable()).to.equal(false);
+  });
+  it('should enable longPolling and disable websockets', () => {
+    forceLongPolling();
+    expect(spyWarn.called).to.equal(false);
+    expect(WebSocketConnection.isAvailable()).to.equal(false);
+    expect(BrowserPollConnection.isAvailable()).to.equal(true);
+  });
+  it('should throw an error when calling forceLongPolling() if TransportManager has already been initialized', () => {
+    TransportManager.globalTransportInitialized_ = true;
+    forceLongPolling();
+    expect(spyWarn).to.have.been.calledWith(transportInitError);
+    expect(WebSocketConnection.isAvailable()).to.equal(false);
+    expect(BrowserPollConnection.isAvailable()).to.equal(true);
+  });
+});