From 08c621ff9c9225c51ee3bc4d9ff40e9ff1f97183 Mon Sep 17 00:00:00 2001
From: Joe <josephnumainville@deephaven.io>
Date: Wed, 28 Feb 2024 16:07:27 -0600
Subject: [PATCH] fix: Added enterprise checks for case where object name is a
 string #36

Fixes #35

Added checks for enterprise info in case where object name is a string

Ran the following code and verified the chart appeared:

from deephaven_ipywidgets import DeephavenWidget

connection_info = "connection"
session_mgr: SessionManager = SessionManager(connection_info)
session_mgr.password("username","password")
session = session_mgr.connect_to_new_worker(name=None, heap_size_gb=1.0)

session.run_script("""
from deephaven import empty_table
from deephaven.plot.figure import Figure
t = empty_table(100).update(["X=i", "Y=i*i"])
f = Figure().plot_xy(series_name="X vs Y", t=t, x="X", y="Y").show()
""")

display(DeephavenWidget("f", session=session))
---
 deephaven_ipywidgets/deephaven.py | 66 ++++++++++++++++++++-----------
 1 file changed, 44 insertions(+), 22 deletions(-)

diff --git a/deephaven_ipywidgets/deephaven.py b/deephaven_ipywidgets/deephaven.py
index 90ff3c9..0aafead 100644
--- a/deephaven_ipywidgets/deephaven.py
+++ b/deephaven_ipywidgets/deephaven.py
@@ -10,6 +10,8 @@
 from __future__ import annotations
 
 import __main__
+from typing import Any
+
 from ipywidgets import DOMWidget
 from traitlets import Unicode, Integer, Bytes, Bool
 from deephaven_server import Server
@@ -19,7 +21,6 @@
 import base64
 import atexit
 
-
 TABLE_TYPES = {"deephaven.table.Table", "pandas.core.frame.DataFrame", "pydeephaven.table.Table"}
 FIGURE_TYPES = {"deephaven.plot.figure.Figure"}
 
@@ -52,6 +53,44 @@ def _cleanup(widget: DeephavenWidget):
     widget.set_trait("kernel_active", False)
 
 
+def _check_session(session: Any, params: dict):
+    """
+    Check the session for a session manager and set the parameters for the widget
+
+    Args:
+        session: The session to check
+        params: The parameters to set
+
+    Returns:
+        str, str: The server URL and token to use
+
+    """
+
+    token = ""
+
+    port = session.port
+    server_url = f"http://{session.host}:{port}/"
+
+    if hasattr(session, "_extra_headers") and b"envoy-prefix" in session._extra_headers:
+        params["envoyPrefix"] = session._extra_headers[b"envoy-prefix"].decode(
+            "ascii"
+        )
+
+    if hasattr(session, "session_manager"):
+        params["authProvider"] = "parent"
+        # We have a DnD session, and we can get the authentication and connection details from the session manager
+        token = base64.b64encode(
+            session.session_manager.auth_client.get_token(
+                "RemoteQueryProcessor"
+            ).SerializeToString()
+        ).decode("us-ascii")
+        server_url = (
+            session.pqinfo().state.connectionDetails.staticUrl
+        )
+
+    return server_url, token
+
+
 class DeephavenWidget(DOMWidget):
     """A wrapper for viewing DeephavenWidgets in IPython"""
 
@@ -96,30 +135,13 @@ def __init__(self, deephaven_object, height=600, width=0, session=None):
                 raise ValueError(
                     "session must be specified when using a remote pydeephaven object by name"
                 )
-            port = session.port
-            server_url = f"http://{session.host}:{port}/"
+
+            server_url, token = _check_session(session, params)
+
         elif _str_object_type(deephaven_object) == "pydeephaven.table.Table":
             session = deephaven_object.session
 
-            if b"envoy-prefix" in session._extra_headers:
-                params["envoyPrefix"] = session._extra_headers[b"envoy-prefix"].decode(
-                    "ascii"
-                )
-
-            port = deephaven_object.session.port
-            server_url = f"http://{deephaven_object.session.host}:{port}/"
-
-            if hasattr(session, "session_manager"):
-                params["authProvider"] = "parent"
-                # We have a DnD session, and we can get the authentication and connection details from the session manager
-                token = base64.b64encode(
-                    session.session_manager.auth_client.get_token(
-                        "RemoteQueryProcessor"
-                    ).SerializeToString()
-                ).decode("us-ascii")
-                server_url = (
-                    deephaven_object.session.pqinfo().state.connectionDetails.staticUrl
-                )
+            server_url, token = _check_session(session, params)
 
             session.bind_table(object_id, deephaven_object)
         else: