From 18158eba0309732c43e7a17ae316ed97073619ec Mon Sep 17 00:00:00 2001 From: vaimdevs Date: Mon, 9 Mar 2026 16:14:57 +0800 Subject: [PATCH 1/6] first try --- graphistry/PlotterBase.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index 28453952ae..95740a928e 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -6,7 +6,7 @@ from graphistry.plugins_types.hypergraph import HypergraphResult from graphistry.render.resolve_render_mode import resolve_render_mode from graphistry.Engine import EngineAbstractType -import copy, hashlib, numpy as np, pandas as pd, pyarrow as pa, sys, uuid +import copy, hashlib, numpy as np, pandas as pd, pyarrow as pa, requests, sys, uuid from functools import lru_cache from weakref import WeakValueDictionary @@ -2238,8 +2238,22 @@ def plot( 'type': 'arrow', 'viztoken': str(uuid.uuid4()) } + url_params = dict(self._url_params) + token = self.api_token() + if token: + try: + server_base = '%s://%s' % (self.session.protocol, self.session.hostname) + resp = requests.post( + '%s/api/v2/auth/jwt/ott/' % server_base, + headers={'Authorization': 'Bearer %s' % token}, + verify=self.session.certificate_validation, + ) + resp.raise_for_status() + url_params['token'] = resp.json()['ott'] + except Exception as e: + logger.warning("Failed to exchange JWT for OTT: %s", e) - viz_url = self._pygraphistry._viz_url(info, self._url_params) + viz_url = self._pygraphistry._viz_url(info, url_params) cfg_client_protocol_hostname = self.session.client_protocol_hostname full_url = ('%s:%s' % (self.session.protocol, viz_url)) if cfg_client_protocol_hostname is None else viz_url From b77b40ddf32c5b19038d59456443a8a734a074f4 Mon Sep 17 00:00:00 2001 From: vaimdevs Date: Mon, 9 Mar 2026 16:26:21 +0800 Subject: [PATCH 2/6] fix --- graphistry/PlotterBase.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index 95740a928e..371f5ecd63 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -2239,7 +2239,7 @@ def plot( 'viztoken': str(uuid.uuid4()) } url_params = dict(self._url_params) - token = self.api_token() + token = self.session.api_token if token: try: server_base = '%s://%s' % (self.session.protocol, self.session.hostname) From 43e2c42b00521cde520a84762c4e700c38311f7d Mon Sep 17 00:00:00 2001 From: vaimdevs Date: Mon, 9 Mar 2026 16:46:15 +0800 Subject: [PATCH 3/6] fix --- graphistry/PlotterBase.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index 371f5ecd63..91ed8d1d0c 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -2244,7 +2244,7 @@ def plot( try: server_base = '%s://%s' % (self.session.protocol, self.session.hostname) resp = requests.post( - '%s/api/v2/auth/jwt/ott/' % server_base, + '%s/api/v1/auth/jwt/ott/' % server_base, headers={'Authorization': 'Bearer %s' % token}, verify=self.session.certificate_validation, ) From 4b913d847fcb95596781f21b1661a06eb1f8a703 Mon Sep 17 00:00:00 2001 From: vaimdevs Date: Wed, 15 Apr 2026 09:13:33 +0800 Subject: [PATCH 4/6] Fix CI --- graphistry/PlotterBase.py | 4 ++-- graphistry/tests/test_trace_headers_behavior.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index 91ed8d1d0c..5c755a098f 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -31,7 +31,7 @@ error, hash_pdf, in_ipython, in_databricks, make_iframe, random_string, warn, cache_coercion, cache_coercion_helper, WeakValueWrapper ) -from graphistry.otel import otel_traced, otel_detail_enabled +from graphistry.otel import otel_traced, otel_detail_enabled, inject_trace_headers from .bolt_util import ( bolt_graph_to_edges_dataframe, @@ -2245,7 +2245,7 @@ def plot( server_base = '%s://%s' % (self.session.protocol, self.session.hostname) resp = requests.post( '%s/api/v1/auth/jwt/ott/' % server_base, - headers={'Authorization': 'Bearer %s' % token}, + headers=inject_trace_headers({'Authorization': 'Bearer %s' % token}), verify=self.session.certificate_validation, ) resp.raise_for_status() diff --git a/graphistry/tests/test_trace_headers_behavior.py b/graphistry/tests/test_trace_headers_behavior.py index 96014e2c0a..1c469542dd 100644 --- a/graphistry/tests/test_trace_headers_behavior.py +++ b/graphistry/tests/test_trace_headers_behavior.py @@ -55,13 +55,17 @@ def _post_response_for_plot(url: str): return _mock_response({"is_valid": True, "is_uploaded": True}) if "/api/v2/share/link/" in url: return _mock_response({"success": True}) + if "/api/v1/auth/jwt/ott/" in url: + return _mock_response({"ott": "test-ott-token"}) raise AssertionError(f"Unexpected POST url: {url}") +@mock.patch("graphistry.PlotterBase.inject_trace_headers") @mock.patch("graphistry.arrow_uploader.inject_trace_headers") @mock.patch("requests.post") -def test_plot_injects_traceparent(mock_post, mock_inject): +def test_plot_injects_traceparent(mock_post, mock_inject, mock_inject_plotter): mock_inject.side_effect = _inject_trace + mock_inject_plotter.side_effect = _inject_trace headers_seen = [] def _fake_post(url, **kwargs): @@ -77,15 +81,17 @@ def _fake_post(url, **kwargs): assert all(h.get("traceparent") == TRACEPARENT for h in headers_seen) +@mock.patch("graphistry.PlotterBase.inject_trace_headers") @mock.patch("graphistry.arrow_uploader.inject_trace_headers") @mock.patch("requests.post") -def test_upload_injects_traceparent(mock_post, mock_inject_uploader): +def test_upload_injects_traceparent(mock_post, mock_inject_uploader, mock_inject_plotter): # Patch ArrowFileUploader module's inject_trace_headers via sys.modules # This is needed because graphistry.ArrowFileUploader resolves to the class, # not the module (due to re-exports in graphistry/__init__.py) arrow_file_uploader_module = sys.modules["graphistry.ArrowFileUploader"] mock_inject_uploader.side_effect = _inject_trace + mock_inject_plotter.side_effect = _inject_trace headers_seen = [] def _fake_post(url, **kwargs): From 5e9a5429c3c3a6e5467e3c2a8e8ad3f472cf77bd Mon Sep 17 00:00:00 2001 From: vaimdevs Date: Wed, 15 Apr 2026 09:32:07 +0800 Subject: [PATCH 5/6] Fix CI --- .../tests/test_trace_headers_behavior.py | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/graphistry/tests/test_trace_headers_behavior.py b/graphistry/tests/test_trace_headers_behavior.py index 1c469542dd..e0633f0d46 100644 --- a/graphistry/tests/test_trace_headers_behavior.py +++ b/graphistry/tests/test_trace_headers_behavior.py @@ -4,9 +4,10 @@ import pandas as pd -# Import the ArrowFileUploader MODULE before graphistry shadows it with the class -# This ensures sys.modules has the module, allowing proper mock patching +# Import modules before graphistry shadows them with classes/symbols. +# This ensures sys.modules has the modules, allowing proper mock patching. import graphistry.ArrowFileUploader as _arrow_file_uploader_module # noqa: F401 +import graphistry.PlotterBase as _plotter_base_module # noqa: F401 import graphistry from graphistry.compute.ast import n, e_forward @@ -60,12 +61,8 @@ def _post_response_for_plot(url: str): raise AssertionError(f"Unexpected POST url: {url}") -@mock.patch("graphistry.PlotterBase.inject_trace_headers") -@mock.patch("graphistry.arrow_uploader.inject_trace_headers") @mock.patch("requests.post") -def test_plot_injects_traceparent(mock_post, mock_inject, mock_inject_plotter): - mock_inject.side_effect = _inject_trace - mock_inject_plotter.side_effect = _inject_trace +def test_plot_injects_traceparent(mock_post): headers_seen = [] def _fake_post(url, **kwargs): @@ -74,24 +71,20 @@ def _fake_post(url, **kwargs): mock_post.side_effect = _fake_post - g = _make_graph() - g.plot(render="g", as_files=False, validate=False, warn=False, memoize=False) + plotter_base_module = sys.modules["graphistry.PlotterBase"] + arrow_uploader_module = sys.modules["graphistry.arrow_uploader"] + + with mock.patch.object(arrow_uploader_module, "inject_trace_headers", side_effect=_inject_trace), \ + mock.patch.object(plotter_base_module, "inject_trace_headers", side_effect=_inject_trace): + g = _make_graph() + g.plot(render="g", as_files=False, validate=False, warn=False, memoize=False) assert headers_seen assert all(h.get("traceparent") == TRACEPARENT for h in headers_seen) -@mock.patch("graphistry.PlotterBase.inject_trace_headers") -@mock.patch("graphistry.arrow_uploader.inject_trace_headers") @mock.patch("requests.post") -def test_upload_injects_traceparent(mock_post, mock_inject_uploader, mock_inject_plotter): - # Patch ArrowFileUploader module's inject_trace_headers via sys.modules - # This is needed because graphistry.ArrowFileUploader resolves to the class, - # not the module (due to re-exports in graphistry/__init__.py) - arrow_file_uploader_module = sys.modules["graphistry.ArrowFileUploader"] - - mock_inject_uploader.side_effect = _inject_trace - mock_inject_plotter.side_effect = _inject_trace +def test_upload_injects_traceparent(mock_post): headers_seen = [] def _fake_post(url, **kwargs): @@ -100,7 +93,17 @@ def _fake_post(url, **kwargs): mock_post.side_effect = _fake_post - with mock.patch.object(arrow_file_uploader_module, "inject_trace_headers", side_effect=_inject_trace): + # Patch inject_trace_headers in all three modules that make POST requests: + # arrow_uploader.py, ArrowFileUploader.py, and PlotterBase.py (OTT exchange). + # Use sys.modules because graphistry/__init__.py re-exports some names as classes, + # shadowing the module attributes on the graphistry package. + arrow_uploader_module = sys.modules["graphistry.arrow_uploader"] + arrow_file_uploader_module = sys.modules["graphistry.ArrowFileUploader"] + plotter_base_module = sys.modules["graphistry.PlotterBase"] + + with mock.patch.object(arrow_uploader_module, "inject_trace_headers", side_effect=_inject_trace), \ + mock.patch.object(arrow_file_uploader_module, "inject_trace_headers", side_effect=_inject_trace), \ + mock.patch.object(plotter_base_module, "inject_trace_headers", side_effect=_inject_trace): g = _make_graph() g.upload(validate=False, warn=False, memoize=False, erase_files_on_fail=False) From 2fcaa234361aa8d8e5bc6c41fd4453fba64fe0dc Mon Sep 17 00:00:00 2001 From: vaimdevs Date: Tue, 28 Apr 2026 15:24:12 +0800 Subject: [PATCH 6/6] Fix --- graphistry/PlotterBase.py | 13 ++++++++----- graphistry/tests/test_trace_headers_behavior.py | 16 ++++++++++++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index 030e2da7eb..4df9703b90 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -2298,7 +2298,10 @@ def plot( 'type': 'arrow', 'viztoken': str(uuid.uuid4()) } - url_params = dict(self._url_params) + # Validate collections in url_params (catches bypass of .collections() method) + from graphistry.validate.validate_collections import normalize_collections_url_params + url_params = normalize_collections_url_params(self._url_params, validate=validate_mode, warn=warn) + token = self.session.api_token if token: try: @@ -2307,16 +2310,16 @@ def plot( '%s/api/v1/auth/jwt/ott/' % server_base, headers=inject_trace_headers({'Authorization': 'Bearer %s' % token}), verify=self.session.certificate_validation, + timeout=30, ) resp.raise_for_status() url_params['token'] = resp.json()['ott'] + except requests.HTTPError as e: + logger.warning("Failed to exchange JWT for OTT: %s (status=%s, body=%.200s)", + e, resp.status_code, resp.text) except Exception as e: logger.warning("Failed to exchange JWT for OTT: %s", e) - # Validate collections in url_params (catches bypass of .collections() method) - from graphistry.validate.validate_collections import normalize_collections_url_params - url_params = normalize_collections_url_params(self._url_params, validate=validate_mode, warn=warn) - viz_url = self._pygraphistry._viz_url(info, url_params) cfg_client_protocol_hostname = self.session.client_protocol_hostname full_url = ('%s:%s' % (self.session.protocol, viz_url)) if cfg_client_protocol_hostname is None else viz_url diff --git a/graphistry/tests/test_trace_headers_behavior.py b/graphistry/tests/test_trace_headers_behavior.py index e0633f0d46..11ef1c1484 100644 --- a/graphistry/tests/test_trace_headers_behavior.py +++ b/graphistry/tests/test_trace_headers_behavior.py @@ -83,6 +83,22 @@ def _fake_post(url, **kwargs): assert all(h.get("traceparent") == TRACEPARENT for h in headers_seen) +@mock.patch("requests.post") +def test_plot_ott_in_url(mock_post): + """OTT from JWT exchange must appear as ?token= in the returned viz URL.""" + mock_post.side_effect = lambda url, **kw: _post_response_for_plot(url) + + plotter_base_module = sys.modules["graphistry.PlotterBase"] + arrow_uploader_module = sys.modules["graphistry.arrow_uploader"] + + with mock.patch.object(arrow_uploader_module, "inject_trace_headers", side_effect=_inject_trace), \ + mock.patch.object(plotter_base_module, "inject_trace_headers", side_effect=_inject_trace): + g = _make_graph() + url = g.plot(render="url", as_files=False, validate=False, warn=False, memoize=False) + + assert "token=test-ott-token" in url, f"OTT missing from viz URL: {url}" + + @mock.patch("requests.post") def test_upload_injects_traceparent(mock_post): headers_seen = []