Skip to content

Commit 7fa2a3e

Browse files
authored
Add support for client feature flags to be passed as query parameters (#6302)
Some plugins (e.g. [image](https://github.com/tensorflow/tensorboard/blob/963afb7be883c15373c121c96c9039828e3fb728/tensorboard/plugins/image/tf_image_dashboard/tf-image-loader.ts#L356) and [audio](https://github.com/tensorflow/tensorboard/blob/963afb7be883c15373c121c96c9039828e3fb728/tensorboard/plugins/audio/tf_audio_dashboard/tf-audio-loader.ts#L111)) specify URLs as src attributes on some HTML elements, which precludes them from attaching the necessary client feature flags via the `requestManager`. This PR adds support to the client feature flags middleware to also accept feature flags via query parameters, so that we can support this use case.
1 parent 136e9e5 commit 7fa2a3e

File tree

3 files changed

+154
-10
lines changed

3 files changed

+154
-10
lines changed

tensorboard/backend/client_feature_flags.py

Lines changed: 60 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""Middleware for injecting client-side feature flags into the Context."""
1616

1717
import json
18+
import urllib.parse
1819

1920
from tensorboard import context
2021
from tensorboard import errors
@@ -24,10 +25,14 @@ class ClientFeatureFlagsMiddleware:
2425
"""Middleware for injecting client-side feature flags into the Context.
2526
2627
The client webapp is expected to include a json-serialized version of its
27-
FeatureFlags in the `X-TensorBoard-Feature-Flags` header. This middleware
28-
extracts the header value and converts it into the client_feature_flags
28+
FeatureFlags in the `X-TensorBoard-Feature-Flags` header or the
29+
`tensorBoardFeatureFlags` query parameter. This middleware extracts the
30+
header or query parameter value and converts it into the client_feature_flags
2931
property for the DataProvider's Context object, where client_feature_flags
3032
is a Dict of string keys and arbitrary value types.
33+
34+
In the event that both the header and query parameter are specified, the
35+
values from the header will take precedence.
3136
"""
3237

3338
def __init__(self, application):
@@ -39,25 +44,70 @@ def __init__(self, application):
3944
self._application = application
4045

4146
def __call__(self, environ, start_response):
42-
possible_feature_flags = environ.get("HTTP_X_TENSORBOARD_FEATURE_FLAGS")
43-
if not possible_feature_flags:
47+
header_feature_flags = self._parse_potential_header_param_flags(
48+
environ.get("HTTP_X_TENSORBOARD_FEATURE_FLAGS")
49+
)
50+
query_string_feature_flags = self._parse_potential_query_param_flags(
51+
environ.get("QUERY_STRING")
52+
)
53+
54+
if not header_feature_flags and not query_string_feature_flags:
4455
return self._application(environ, start_response)
4556

57+
# header flags take precedence
58+
for flag, value in header_feature_flags.items():
59+
query_string_feature_flags[flag] = value
60+
61+
ctx = context.from_environ(environ).replace(
62+
client_feature_flags=query_string_feature_flags
63+
)
64+
context.set_in_environ(environ, ctx)
65+
66+
return self._application(environ, start_response)
67+
68+
def _parse_potential_header_param_flags(self, header_string):
69+
if not header_string:
70+
return {}
71+
4672
try:
47-
client_feature_flags = json.loads(possible_feature_flags)
73+
header_feature_flags = json.loads(header_string)
4874
except json.JSONDecodeError:
4975
raise errors.InvalidArgumentError(
5076
"X-TensorBoard-Feature-Flags cannot be JSON decoded."
5177
)
5278

53-
if not isinstance(client_feature_flags, dict):
79+
if not isinstance(header_feature_flags, dict):
5480
raise errors.InvalidArgumentError(
5581
"X-TensorBoard-Feature-Flags cannot be decoded to a dict."
5682
)
5783

58-
ctx = context.from_environ(environ).replace(
59-
client_feature_flags=client_feature_flags
84+
return header_feature_flags
85+
86+
def _parse_potential_query_param_flags(self, query_string):
87+
if not query_string:
88+
return {}
89+
90+
try:
91+
query_string_json = urllib.parse.parse_qs(query_string)
92+
except ValueError:
93+
return {}
94+
95+
# parse_qs returns the dictionary values as lists for each name.
96+
potential_feature_flags = query_string_json.get(
97+
"tensorBoardFeatureFlags", []
6098
)
61-
context.set_in_environ(environ, ctx)
99+
if not potential_feature_flags:
100+
return {}
101+
try:
102+
client_feature_flags = json.loads(potential_feature_flags[0])
103+
except json.JSONDecodeError:
104+
raise errors.InvalidArgumentError(
105+
"tensorBoardFeatureFlags cannot be JSON decoded."
106+
)
62107

63-
return self._application(environ, start_response)
108+
if not isinstance(client_feature_flags, dict):
109+
raise errors.InvalidArgumentError(
110+
"tensorBoardFeatureFlags cannot be decoded to a dict."
111+
)
112+
113+
return client_feature_flags

tensorboard/backend/client_feature_flags_test.py

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,13 +52,27 @@ def test_no_header(self):
5252
response = server.get("")
5353
self._assert_ok(response, {})
5454

55+
def test_no_query_string(self):
56+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
57+
server = werkzeug_test.Client(app, wrappers.Response)
58+
59+
response = server.get("")
60+
self._assert_ok(response, {})
61+
5562
def test_header_with_no_value(self):
5663
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
5764
server = werkzeug_test.Client(app, wrappers.Response)
5865

5966
response = server.get("", headers=[("X-TensorBoard-Feature-Flags", "")])
6067
self._assert_ok(response, {})
6168

69+
def test_query_string_with_no_value(self):
70+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
71+
server = werkzeug_test.Client(app, wrappers.Response)
72+
73+
response = server.get("", query_string={"tensorBoardFeatureFlags": ""})
74+
self._assert_ok(response, {})
75+
6276
def test_header_with_no_flags(self):
6377
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
6478
server = werkzeug_test.Client(app, wrappers.Response)
@@ -68,6 +82,15 @@ def test_header_with_no_flags(self):
6882
)
6983
self._assert_ok(response, {})
7084

85+
def test_query_string_with_no_flags(self):
86+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
87+
server = werkzeug_test.Client(app, wrappers.Response)
88+
89+
response = server.get(
90+
"", query_string={"tensorBoardFeatureFlags": "{}"}
91+
)
92+
self._assert_ok(response, {})
93+
7194
def test_header_with_client_feature_flags(self):
7295
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
7396
server = werkzeug_test.Client(app, wrappers.Response)
@@ -90,6 +113,25 @@ def test_header_with_client_feature_flags(self):
90113
},
91114
)
92115

116+
def test_query_string_with_client_feature_flags(self):
117+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
118+
server = werkzeug_test.Client(app, wrappers.Response)
119+
120+
response = server.get(
121+
"",
122+
query_string={
123+
"tensorBoardFeatureFlags": '{"str": "hi", "bool": true, "strArr": ["one", "two"]}'
124+
},
125+
)
126+
self._assert_ok(
127+
response,
128+
{
129+
"str": "hi",
130+
"bool": True,
131+
"strArr": ["one", "two"],
132+
},
133+
)
134+
93135
def test_header_with_json_not_decodable(self):
94136
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
95137
server = werkzeug_test.Client(app, wrappers.Response)
@@ -107,6 +149,20 @@ def test_header_with_json_not_decodable(self):
107149
],
108150
)
109151

152+
def test_query_string_with_json_not_decodable(self):
153+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
154+
server = werkzeug_test.Client(app, wrappers.Response)
155+
156+
with self.assertRaisesRegex(
157+
errors.InvalidArgumentError, "cannot be JSON decoded."
158+
):
159+
response = server.get(
160+
"",
161+
query_string={
162+
"tensorBoardFeatureFlags": "some_invalid_json {} {}",
163+
},
164+
)
165+
110166
def test_header_with_json_not_dict(self):
111167
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
112168
server = werkzeug_test.Client(app, wrappers.Response)
@@ -124,6 +180,43 @@ def test_header_with_json_not_dict(self):
124180
],
125181
)
126182

183+
def test_query_string_with_json_not_dict(self):
184+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
185+
server = werkzeug_test.Client(app, wrappers.Response)
186+
187+
with self.assertRaisesRegex(
188+
errors.InvalidArgumentError, "cannot be decoded to a dict"
189+
):
190+
response = server.get(
191+
"",
192+
query_string={
193+
"tensorBoardFeatureFlags": '["not", "a", "dict"]',
194+
},
195+
)
196+
197+
def test_header_feature_flags_take_precedence(self):
198+
app = client_feature_flags.ClientFeatureFlagsMiddleware(self._echo_app)
199+
server = werkzeug_test.Client(app, wrappers.Response)
200+
201+
response = server.get(
202+
"",
203+
headers=[
204+
(
205+
"X-TensorBoard-Feature-Flags",
206+
'{"a": "1", "b": "2"}',
207+
)
208+
],
209+
query_string={"tensorBoardFeatureFlags": '{"a": "2", "c": "3"}'},
210+
)
211+
self._assert_ok(
212+
response,
213+
{
214+
"a": "1",
215+
"b": "2",
216+
"c": "3",
217+
},
218+
)
219+
127220

128221
if __name__ == "__main__":
129222
tb_test.main()

tensorboard/webapp/feature_flag/http/const.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,4 @@ limitations under the License.
1414
==============================================================================*/
1515

1616
export const FEATURE_FLAGS_HEADER_NAME = 'X-TensorBoard-Feature-Flags';
17+
export const FEATURE_FLAGS_QUERY_STRING_NAME = 'tensorBoardFeatureFlags';

0 commit comments

Comments
 (0)