Add directives to modify attributes and include feature flags where n…

…eeded (#6341)

Updates individual HTML elements to include client feature flags as
query parameters on their src attributes through the use of a directive.

Validated via a local tensorboard that the attributes generated from
these directives include the client feature flags as query parameters.

tensorboard/webapp/feature_flag/BUILD

@@ -43,6 +43,7 @@ tf_ts_library(
 tf_ng_web_test_suite(
     name = "karma_test",
     deps = [
+        "//tensorboard/webapp/feature_flag/directives:directives_test_lib",
         "//tensorboard/webapp/feature_flag/effects:effects_test_lib",
         "//tensorboard/webapp/feature_flag/http:http_test_lib",
         "//tensorboard/webapp/feature_flag/store:store_test_lib",

tensorboard/webapp/feature_flag/directives/BUILD

@@ -0,0 +1,45 @@
+load("//tensorboard/defs:defs.bzl", "tf_ng_module", "tf_ts_library")
+
+package(default_visibility = ["//tensorboard:internal"])
+
+tf_ng_module(
+    name = "directives",
+    srcs = [
+        "feature_flag_directive.ts",
+        "feature_flag_directive_module.ts",
+    ],
+    deps = [
+        "//tensorboard/webapp/angular:expect_angular_common_http",
+        "//tensorboard/webapp/feature_flag:types",
+        "//tensorboard/webapp/feature_flag/http:const",
+        "//tensorboard/webapp/feature_flag/store",
+        "//tensorboard/webapp/feature_flag/store:types",
+        "@npm//@angular/core",
+        "@npm//@ngrx/store",
+        "@npm//rxjs",
+    ],
+)
+
+tf_ts_library(
+    name = "directives_test_lib",
+    testonly = True,
+    srcs = [
+        "feature_flag_directive_test.ts",
+    ],
+    deps = [
+        ":directives",
+        "//tensorboard/webapp:app_state",
+        "//tensorboard/webapp/feature_flag:testing",
+        "//tensorboard/webapp/feature_flag:types",
+        "//tensorboard/webapp/feature_flag/http:const",
+        "//tensorboard/webapp/feature_flag/store",
+        "//tensorboard/webapp/feature_flag/store:types",
+        "//tensorboard/webapp/testing:utils",
+        "@npm//@angular/core",
+        "@npm//@angular/platform-browser",
+        "@npm//@ngrx/effects",
+        "@npm//@ngrx/store",
+        "@npm//@types/jasmine",
+        "@npm//rxjs",
+    ],
+)

tensorboard/webapp/feature_flag/directives/feature_flag_directive.ts

@@ -0,0 +1,68 @@
+/* Copyright 2023 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 {Directive, ElementRef, HostBinding, Input} from '@angular/core';
+import {Store} from '@ngrx/store';
+import {firstValueFrom} from 'rxjs';
+import {map, tap} from 'rxjs/operators';
+
+import {FEATURE_FLAGS_QUERY_STRING_NAME} from '../http/const';
+import {getFeatureFlagsToSendToServer} from '../store/feature_flag_selectors';
+import {State as FeatureFlagState} from '../store/feature_flag_types';
+
+@Directive({selector: 'a[includeFeatureFlags], img[includeFeatureFlags]'})
+export class FeatureFlagDirective {
+  @HostBinding('attr.href') hrefAttr: string | undefined = undefined;
+  @HostBinding('attr.src') srcAttr: string | undefined = undefined;
+  // The selector applies if [includeFeatureFlags] is present at all. Supplying
+  // [includeFeatureFlags]="true"/"false" has no impact on the actual logic and
+  // will behave the same as though [includeFeatureFlags] is unset.
+  @Input() includeFeatureFlags: boolean = true;
+
+  constructor(private readonly store: Store<FeatureFlagState>) {}
+
+  private getUrlWithFeatureFlags(url: string) {
+    return this.store.select(getFeatureFlagsToSendToServer).pipe(
+      map((featureFlags) => {
+        if (Object.keys(featureFlags).length > 0) {
+          const params = new URLSearchParams([
+            [FEATURE_FLAGS_QUERY_STRING_NAME, JSON.stringify(featureFlags)],
+          ]);
+          const delimiter = url.includes('?') ? '&' : '?';
+          return url + delimiter + String(params);
+        } else {
+          return url;
+        }
+      })
+    );
+  }
+
+  @Input()
+  set href(href: string | null) {
+    if (href) {
+      firstValueFrom(this.getUrlWithFeatureFlags(href)).then((value) => {
+        this.hrefAttr = value;
+      });
+    }
+  }
+
+  @Input()
+  set src(src: string | null) {
+    if (src) {
+      firstValueFrom(this.getUrlWithFeatureFlags(src)).then((value) => {
+        this.srcAttr = value;
+      });
+    }
+  }
+}

tensorboard/webapp/feature_flag/directives/feature_flag_directive_module.ts

@@ -0,0 +1,22 @@
+/* Copyright 2023 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 {NgModule} from '@angular/core';
+import {FeatureFlagDirective} from './feature_flag_directive';
+
+@NgModule({
+  declarations: [FeatureFlagDirective],
+  exports: [FeatureFlagDirective],
+})
+export class FeatureFlagDirectiveModule {}

tensorboard/webapp/feature_flag/directives/feature_flag_directive_test.ts

@@ -0,0 +1,140 @@
+/* Copyright 2023 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 {Component, DebugElement, Input} from '@angular/core';
+import {TestBed, fakeAsync, tick} from '@angular/core/testing';
+import {By} from '@angular/platform-browser';
+import {Store} from '@ngrx/store';
+import {MockStore} from '@ngrx/store/testing';
+
+import {provideMockTbStore} from '../../testing/utils';
+import {FEATURE_FLAGS_HEADER_NAME} from '../http/const';
+import {getFeatureFlagsToSendToServer} from '../store/feature_flag_selectors';
+import {State as FeatureFlagState} from '../store/feature_flag_types';
+
+import {FeatureFlagDirective} from './feature_flag_directive';
+
+@Component({
+  selector: 'test-matching-selector',
+  template: `
+    <p>
+      <a [href]="href" [includeFeatureFlags]>test link</a>
+      <img [src]="src" [includeFeatureFlags] />
+    </p>
+  `,
+})
+export class TestMatchingComponent {
+  @Input() href!: string;
+  @Input() src!: string;
+}
+
+@Component({
+  selector: 'test-nonmatching-selector',
+  template: `
+    <p>
+      <img [src]="src" />
+    </p>
+  `,
+})
+export class TestNonmatchingComponent {
+  @Input() src!: string;
+}
+
+describe('feature_flags', () => {
+  let store: MockStore<FeatureFlagState>;
+
+  beforeEach(async () => {
+    await TestBed.configureTestingModule({
+      providers: [provideMockTbStore()],
+      declarations: [
+        TestMatchingComponent,
+        TestNonmatchingComponent,
+        FeatureFlagDirective,
+      ],
+    }).compileComponents();
+
+    store = TestBed.inject<Store<FeatureFlagState>>(
+      Store
+    ) as MockStore<FeatureFlagState>;
+    store.overrideSelector(getFeatureFlagsToSendToServer, {});
+  });
+
+  afterEach(() => {
+    store?.resetSelectors();
+  });
+
+  function createMatchingHrefComponent(href: string): DebugElement {
+    const fixture = TestBed.createComponent(TestMatchingComponent);
+    fixture.componentInstance.href = href;
+    fixture.detectChanges();
+    tick();
+    fixture.detectChanges();
+    return fixture.debugElement.query(By.css('a'));
+  }
+
+  function createMatchingImgComponent(src: string): DebugElement {
+    const fixture = TestBed.createComponent(TestMatchingComponent);
+    fixture.componentInstance.src = src;
+    fixture.detectChanges();
+    tick();
+    fixture.detectChanges();
+    return fixture.debugElement.query(By.css('img'));
+  }
+
+  function createNonmatchingImgComponent(src: string): DebugElement {
+    const fixture = TestBed.createComponent(TestNonmatchingComponent);
+    fixture.componentInstance.src = src;
+    fixture.detectChanges();
+    tick();
+    fixture.detectChanges();
+    return fixture.debugElement.query(By.css('img'));
+  }
+
+  it('injects feature flags in <img> tags if any are set without preexisting query parameters', fakeAsync(() => {
+    store.overrideSelector(getFeatureFlagsToSendToServer, {inColab: true});
+    const anchorStr = createMatchingImgComponent('https://abc.def');
+    expect(anchorStr.attributes['src']).toBe(
+      'https://abc.def?tensorBoardFeatureFlags=%7B%22inColab%22%3Atrue%7D'
+    );
+  }));
+
+  it('injects feature flags in <img> tags if any are set with preexisting query parameters', fakeAsync(() => {
+    store.overrideSelector(getFeatureFlagsToSendToServer, {inColab: true});
+    const anchorStr = createMatchingImgComponent('https://abc.def?test=value');
+    expect(anchorStr.attributes['src']).toBe(
+      'https://abc.def?test=value&tensorBoardFeatureFlags=%7B%22inColab%22%3Atrue%7D'
+    );
+  }));
+
+  it('leaves <img> tags unmodified if no feature flags are set', fakeAsync(() => {
+    const anchorStr = createMatchingImgComponent('https://abc.def');
+    expect(anchorStr.attributes['src']).toBe('https://abc.def');
+  }));
+
+  it('leaves <img> tags unmodified if [includeFeatureFlags] is not included', fakeAsync(() => {
+    store.overrideSelector(getFeatureFlagsToSendToServer, {inColab: true});
+    const anchorStr = createNonmatchingImgComponent(
+      'https://abc.def?test=value'
+    );
+    expect(anchorStr.attributes['src']).toBe('https://abc.def?test=value');
+  }));
+
+  it('injects feature flags in <a> tags if any are set', fakeAsync(() => {
+    store.overrideSelector(getFeatureFlagsToSendToServer, {inColab: true});
+    const anchorStr = createMatchingHrefComponent('https://abc.def');
+    expect(anchorStr.attributes['href']).toBe(
+      'https://abc.def?tensorBoardFeatureFlags=%7B%22inColab%22%3Atrue%7D'
+    );
+  }));
+});

tensorboard/webapp/metrics/views/card_renderer/BUILD

@@ -65,6 +65,7 @@ tf_ng_module(
         "//tensorboard/webapp/angular:expect_angular_material_dialog",
         "//tensorboard/webapp/angular:expect_angular_material_input",
         "//tensorboard/webapp/angular:expect_angular_material_select",
+        "//tensorboard/webapp/feature_flag/directives",
         "//tensorboard/webapp/metrics:types",
         "//tensorboard/webapp/metrics/data_source",
         "//tensorboard/webapp/metrics/data_source:types",
@@ -158,6 +159,7 @@ tf_ng_module(
         "//tensorboard/webapp/angular:expect_angular_material_icon",
         "//tensorboard/webapp/angular:expect_angular_material_progress_spinner",
         "//tensorboard/webapp/angular:expect_angular_material_slider",
+        "//tensorboard/webapp/feature_flag/directives",
         "//tensorboard/webapp/metrics:types",
         "//tensorboard/webapp/metrics/actions",
         "//tensorboard/webapp/metrics/data_source",

tensorboard/webapp/metrics/views/card_renderer/data_download_dialog_component.ng.html

@@ -54,14 +54,16 @@ <h2>
         mat-stroked-button
         [disabled]="!downloadUrlJson"
         [download]="getDownloadName('json')"
-        [attr.href]="downloadUrlJson"
+        [href]="downloadUrlJson"
+        [includeFeatureFlags]
         >JSON</a
       >
       <a
         mat-stroked-button
         [disabled]="!downloadUrlCsv"
         [download]="getDownloadName('csv')"
-        [attr.href]="downloadUrlCsv"
+        [href]="downloadUrlCsv"
+        [includeFeatureFlags]
         >CSV</a
       >
     </div>

tensorboard/webapp/metrics/views/card_renderer/data_download_dialog_test.ts

@@ -258,7 +258,9 @@ describe('metrics/views/data_download_dialog', () => {
     const pathAndQueries = downloadLinks.map((link) => {
       return link.nativeElement.href.slice(link.nativeElement.origin.length);
     });
-    expect(pathAndQueries).toEqual(['', '']);
+    // Karma seems to inject this value when using [href] instead of
+    // [attr.href] in the template.
+    expect(pathAndQueries).toEqual(['/context.html', '/context.html']);
   });
 
   it('forms correct url even if run names are all the same', async () => {

tensorboard/webapp/metrics/views/card_renderer/data_download_module.ts

@@ -19,6 +19,7 @@ import {MatButtonModule} from '@angular/material/button';
 import {MatDialogModule} from '@angular/material/dialog';
 import {MatInputModule} from '@angular/material/input';
 import {MatSelectModule} from '@angular/material/select';
+import {FeatureFlagDirectiveModule} from '../../../feature_flag/directives/feature_flag_directive_module';
 import {MetricsDataSourceModule} from '../../data_source';
 import {DataDownloadDialogComponent} from './data_download_dialog_component';
 import {DataDownloadDialogContainer} from './data_download_dialog_container';
@@ -28,6 +29,7 @@ import {DataDownloadDialogContainer} from './data_download_dialog_container';
   exports: [DataDownloadDialogContainer],
   imports: [
     CommonModule,
+    FeatureFlagDirectiveModule,
     FormsModule,
     MatButtonModule,
     MatDialogModule,

tensorboard/webapp/metrics/views/card_renderer/image_card_component.ng.html

@@ -113,6 +113,7 @@
       alt="Image at step {{ steps[stepIndex] }}"
       src="{{ imageUrl }}"
       [ngStyle]="{'filter': cssFilter()}"
+      [includeFeatureFlags]
     />
   </div>
 </ng-container>

tensorboard/webapp/metrics/views/card_renderer/image_card_module.ts

@@ -21,6 +21,7 @@ import {MatSliderModule} from '@angular/material/slider';
 import {TruncatedPathModule} from '../../../widgets/text/truncated_path_module';
 import {ImageCardComponent} from './image_card_component';
 import {ImageCardContainer} from './image_card_container';
+import {FeatureFlagDirectiveModule} from '../../../feature_flag/directives/feature_flag_directive_module';
 import {RunNameModule} from './run_name_module';
 import {VisLinkedTimeSelectionWarningModule} from './vis_linked_time_selection_warning_module';
 
@@ -29,6 +30,7 @@ import {VisLinkedTimeSelectionWarningModule} from './vis_linked_time_selection_w
   exports: [ImageCardContainer],
   imports: [
     CommonModule,
+    FeatureFlagDirectiveModule,
     MatButtonModule,
     MatIconModule,
     MatProgressSpinnerModule,