Adds mechanism to include feature flags in Polymer HTTP requests

tensorboard/components/BUILD

@@ -30,6 +30,7 @@ tf_ts_library(
     deps = [
         "//tensorboard/components/tf_backend",
         "//tensorboard/components/tf_color_scale",
+        "//tensorboard/components/tf_feature_flags",
         "//tensorboard/components/tf_globals",
         "//tensorboard/components/tf_markdown_view",
         "//tensorboard/components/tf_paginated_view",

tensorboard/components/polymer3_interop_helper.ts

@@ -18,6 +18,7 @@ limitations under the License.
  */
 import './tf_backend/tf-backend-polymer';
 import './tf_color_scale/tf-color-scale-polymer';
+import './tf_feature_flags/tf-feature-flags-polymer';
 import './tf_globals/globals-polymer';
 import './tf_paginated_view/tf-paginated-view-store';
 import './tf_storage/tf-storage-polymer';

tensorboard/components/tf_backend/BUILD

@@ -22,7 +22,9 @@ tf_ts_library(
     strict_checks = False,
     deps = [
         ":type",
+        "//tensorboard/components/tf_feature_flags",
         "//tensorboard/components/vz_sorting",
+        "//tensorboard/webapp/feature_flag/http:const",
         "@npm//@polymer/decorators",
         "@npm//@polymer/polymer",
         "@npm//@types/lodash",

tensorboard/components/tf_backend/requestManager.ts

@@ -12,6 +12,9 @@ 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 {FEATURE_FLAGS_HEADER_NAME} from '../../webapp/feature_flag/http/const';
+import {getFeatureFlags} from '../tf_feature_flags/feature-flags';
+
 interface ResolveReject {
   resolve: (x: unknown) => void;
   reject: (x: unknown) => void;
@@ -240,6 +243,10 @@ export class RequestManager {
         requestOptions.withCredentials,
         requestOptions.contentType
       );
+      req.setRequestHeader(
+        FEATURE_FLAGS_HEADER_NAME,
+        JSON.stringify(getFeatureFlags())
+      );
       req.onload = function () {
         if (req.status === 200) {
           resolve(JSON.parse(req.responseText) as any);

tensorboard/components/tf_feature_flags/BUILD

@@ -0,0 +1,38 @@
+load("//tensorboard/defs:defs.bzl", "tf_ng_web_test_suite", "tf_ts_library")
+
+package(default_visibility = ["//tensorboard:internal"])
+
+licenses(["notice"])
+
+tf_ts_library(
+    name = "tf_feature_flags",
+    srcs = [
+        "feature-flags.ts",
+        "tf-feature-flags-polymer.ts",
+    ],
+    deps = [
+        "//tensorboard/webapp/feature_flag:types",
+        "@npm//@polymer/decorators",
+        "@npm//@polymer/polymer",
+    ],
+)
+
+tf_ts_library(
+    name = "test_lib",
+    testonly = True,
+    srcs = [
+        "feature-flags-test.ts",
+    ],
+    deps = [
+        ":tf_feature_flags",
+        "//tensorboard/webapp/feature_flag:testing",
+        "@npm//@types/jasmine",
+    ],
+)
+
+tf_ng_web_test_suite(
+    name = "karma_test",
+    deps = [
+        ":test_lib",
+    ],
+)

tensorboard/components/tf_feature_flags/feature-flags-test.ts

@@ -0,0 +1,36 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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 {buildFeatureFlag} from '../../webapp/feature_flag/testing';
+import * as feature_flags from './feature-flags';
+
+describe('feature-flags', () => {
+  beforeEach(() => {
+    feature_flags.initializeFeatureFlags();
+  });
+
+  it('sets and gets FeatureFlags', () => {
+    feature_flags.setFeatureFlags(buildFeatureFlag({inColab: true}));
+    expect(feature_flags.getFeatureFlags()).toEqual(
+      buildFeatureFlag({inColab: true})
+    );
+  });
+
+  it('throws Error if getFeatureFlags is called before setFeatureFlags', () => {
+    expect(() => feature_flags.getFeatureFlags()).toThrow(
+      new Error('FeatureFlags have not yet been determined by TensorBoard.')
+    );
+  });
+});

tensorboard/components/tf_feature_flags/feature-flags.ts

@@ -0,0 +1,50 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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 {FeatureFlags} from '../../webapp/feature_flag/types';
+
+// FeatureFlags are managed and injected by the Angular portion of the TB code
+// base. In practice they are set one time soon after application load and never
+// again for the lifetime of the application.
+let _featureFlags: FeatureFlags | null;
+initializeFeatureFlags();
+
+export function initializeFeatureFlags(): void {
+  _featureFlags = null;
+}
+
+/**
+ * Sets the FeatureFlags for use in the Polymer portion of the TB code base.
+ *
+ * In practice this should only be called by the Angular portion of the TB code
+ * base immediately after it determines the final set of FeatureFlags.
+ */
+export function setFeatureFlags(featureFlags: FeatureFlags): void {
+  _featureFlags = featureFlags;
+}
+
+/**
+ * Retrieves FeatureFlags.
+ *
+ * @throws Error if FeatureFlags have not yet been set. In practice they should
+ *     be set soon after application load before any Polymer code is invoked.
+ *     This runtime check acts as a sanity check to enforce that assumption.
+ */
+export function getFeatureFlags(): FeatureFlags {
+  if (_featureFlags === null) {
+    throw Error('FeatureFlags have not yet been determined by TensorBoard.');
+  }
+  return _featureFlags;
+}

tensorboard/components/tf_feature_flags/tf-feature-flags-polymer.ts

@@ -0,0 +1,27 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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 {customElement} from '@polymer/decorators';
+import {PolymerElement} from '@polymer/polymer';
+import * as tf_feature_flags from './feature-flags';
+
+/**
+ * This Polymer component allows the underlying API to be accessible from
+ * Angular TensorBoard by exposing otherwise mangled smybols.
+ */
+@customElement('tf-feature-flags')
+class TfFeatureFlags extends PolymerElement {
+  override _template = null;
+  tf_feature_flags = tf_feature_flags;
+}

tensorboard/webapp/BUILD

@@ -79,6 +79,9 @@ tf_ts_library(
     srcs = [
         "tb_polymer_interop_types.ts",
     ],
+    deps = [
+        "//tensorboard/webapp/feature_flag:types",
+    ],
 )
 
 tf_sass_binary(

tensorboard/webapp/feature_flag/effects/BUILD

@@ -8,6 +8,7 @@ tf_ng_module(
         "feature_flag_effects.ts",
     ],
     deps = [
+        "//tensorboard/webapp:tb_polymer_interop_types",
         "//tensorboard/webapp/feature_flag:force_svg_data_source",
         "//tensorboard/webapp/feature_flag/actions",
         "//tensorboard/webapp/feature_flag/store",

tensorboard/webapp/feature_flag/effects/feature_flag_effects.ts

@@ -16,17 +16,27 @@ limitations under the License.
 import {Injectable} from '@angular/core';
 import {Actions, createEffect, ofType} from '@ngrx/effects';
 import {Action, createAction, Store} from '@ngrx/store';
-import {combineLatestWith, map} from 'rxjs/operators';
+import {combineLatestWith, map, tap, withLatestFrom} from 'rxjs/operators';
+import '../../tb_polymer_interop_types';
 import {TBFeatureFlagDataSource} from '../../webapp_data_source/tb_feature_flag_data_source_types';
 import {partialFeatureFlagsLoaded} from '../actions/feature_flag_actions';
 import {ForceSvgDataSource} from '../force_svg_data_source';
-import {getIsAutoDarkModeAllowed} from '../store/feature_flag_selectors';
+import {
+  getFeatureFlags,
+  getIsAutoDarkModeAllowed,
+} from '../store/feature_flag_selectors';
 import {State} from '../store/feature_flag_types';
 
 const effectsInitialized = createAction('[FEATURE FLAG] Effects Init');
 
 @Injectable()
 export class FeatureFlagEffects {
+  // Ngrx assumes all Effect classes have properties that inherit from the base
+  // JS Object. `tf_feature_flags` does not, so we wrap it.
+  private readonly tfFeatureFlags = {
+    ref: document.createElement('tf-feature-flags').tf_feature_flags,
+  };
+
   /** @export */
   readonly getFeatureFlags$ = createEffect(() =>
     this.actions$.pipe(
@@ -46,6 +56,27 @@ export class FeatureFlagEffects {
     )
   );
 
+  /**
+   * Pass FeatureFlags to the Polymer portion of the code base immediately after
+   * feature flags have been finalized.
+   *
+   * @export
+   */
+  readonly updatePolymerFeatureFlags$ = createEffect(
+    () =>
+      this.actions$.pipe(
+        // partialFeatureFlagsLoaded triggers this effect but the actual
+        // feature flag values used are from the Store, given that it contains
+        // the finalized merged feature flags.
+        ofType(partialFeatureFlagsLoaded),
+        withLatestFrom(this.store.select(getFeatureFlags)),
+        tap(([, featureFlags]) => {
+          this.tfFeatureFlags.ref.setFeatureFlags(featureFlags);
+        })
+      ),
+    {dispatch: false}
+  );
+
   constructor(
     private readonly actions$: Actions,
     private readonly store: Store<State>,

tensorboard/webapp/feature_flag/effects/feature_flag_effects_test.ts

@@ -25,7 +25,10 @@ import {
 import {partialFeatureFlagsLoaded} from '../actions/feature_flag_actions';
 import {ForceSvgDataSource} from '../force_svg_data_source';
 import {ForceSvgDataSourceModule} from '../force_svg_data_source_module';
-import {getIsAutoDarkModeAllowed} from '../store/feature_flag_selectors';
+import {
+  getFeatureFlags,
+  getIsAutoDarkModeAllowed,
+} from '../store/feature_flag_selectors';
 import {State} from '../store/feature_flag_types';
 import {buildFeatureFlag} from '../testing';
 import {FeatureFlags} from '../types';
@@ -37,6 +40,7 @@ describe('feature_flag_effects', () => {
   let dataSource: TestingTBFeatureFlagDataSource;
   let forceSvgDataSource: ForceSvgDataSource;
   let effects: FeatureFlagEffects;
+  let setPolymerFeatureFlagsSpy: jasmine.Spy;
 
   beforeEach(async () => {
     actions = new ReplaySubject<Action>(1);
@@ -48,6 +52,16 @@ describe('feature_flag_effects', () => {
         provideMockStore(),
       ],
     }).compileComponents();
+
+    setPolymerFeatureFlagsSpy = jasmine.createSpy('setFeatureFlags');
+    const createElementSpy = spyOn(document, 'createElement');
+    createElementSpy.withArgs('tf-feature-flags').and.returnValue({
+      tf_feature_flags: {
+        setFeatureFlags: setPolymerFeatureFlagsSpy,
+      },
+    } as unknown as HTMLElement);
+    createElementSpy.and.callThrough();
+
     effects = TestBed.inject(FeatureFlagEffects);
     store = TestBed.inject<Store<State>>(Store) as MockStore<State>;
     dataSource = TestBed.inject(TestingTBFeatureFlagDataSource);
@@ -135,4 +149,32 @@ describe('feature_flag_effects', () => {
       expect(updateSpy).toHaveBeenCalledTimes(0);
     });
   });
+
+  describe('updatePolymerFeatureFlags$', () => {
+    it('sets polymer feature flags after data source fetch', () => {
+      // This represents the complete FeatureFlags object as calculated by the
+      // Store after the feature flags data source fetch.
+      store.overrideSelector(
+        getFeatureFlags,
+        buildFeatureFlag({inColab: true})
+      );
+      store.refreshState();
+
+      effects.updatePolymerFeatureFlags$.subscribe();
+
+      actions.next(
+        partialFeatureFlagsLoaded({
+          // This represents the incomplete FeatureFlags object that has just
+          // been fetched by the feature flags data source.
+          features: buildFeatureFlag({inColab: false}),
+        })
+      );
+
+      // Expect tf_feature_flags.setFeatureFlags() to be called using the
+      // FeatureFlags object from the Store (and not from the action).
+      expect(setPolymerFeatureFlagsSpy).toHaveBeenCalledOnceWith(
+        buildFeatureFlag({inColab: true})
+      );
+    });
+  });
 });

tensorboard/webapp/feature_flag/http/BUILD

@@ -2,12 +2,20 @@ load("//tensorboard/defs:defs.bzl", "tf_ng_module", "tf_ts_library")
 
 package(default_visibility = ["//tensorboard:internal"])
 
+tf_ng_module(
+    name = "const",
+    srcs = [
+        "const.ts",
+    ],
+)
+
 tf_ng_module(
     name = "http",
     srcs = [
         "feature_flag_http_interceptor.ts",
     ],
     deps = [
+        ":const",
         "//tensorboard/webapp/angular:expect_angular_common_http",
         "//tensorboard/webapp/feature_flag/store",
         "//tensorboard/webapp/feature_flag/store:types",
@@ -24,6 +32,7 @@ tf_ts_library(
         "feature_flag_http_interceptor_test.ts",
     ],
     deps = [
+        ":const",
         ":http",
         "//tensorboard/webapp/angular:expect_angular_common_http",
         "//tensorboard/webapp/angular:expect_angular_common_http_testing",

tensorboard/webapp/feature_flag/http/const.ts

@@ -0,0 +1,16 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+export const FEATURE_FLAGS_HEADER_NAME = 'X-TensorBoard-Feature-Flags';