Resolve the discovery service and use it as the base URL for the refresh component

- Use the resolved URL when initializing the refresh component
- Perform the resolution in PanelClient and expose it as a property

Signed-off-by: Joe Friedrichsen <[email protected]>

Author: jfriedri-ni

src/nipanel/_panel_client.py

@@ -15,7 +15,7 @@
     SetValueRequest,
 )
 from ni.pythonpanel.v1.python_panel_service_pb2_grpc import PythonPanelServiceStub
-from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
+from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient, ServiceLocation
 from ni_measurement_plugin_sdk_service.grpc.channelpool import GrpcChannelPool
 from typing_extensions import ParamSpec
 
@@ -57,8 +57,14 @@ def __init__(
         self._discovery_client = discovery_client
         self._grpc_channel_pool = grpc_channel_pool
         self._grpc_channel = grpc_channel
+        self._proxy_location: ServiceLocation | None = None
         self._stub: PythonPanelServiceStub | None = None
 
+    @property
+    def proxy_location(self) -> ServiceLocation:
+        """Return the ServiceLocation for the panel's proxy."""
+        return self._get_proxy_location()
+
     def start_panel(self, panel_id: str, panel_script_path: str, python_path: str) -> str:
         """Start the panel.
 
@@ -134,29 +140,61 @@ def get_value(self, panel_id: str, value_id: str) -> tuple[bool, object]:
         else:
             return False, None
 
+    def _get_proxy_location(self) -> ServiceLocation:
+        if self._proxy_location is None:
+            self._proxy_location = self._resolve_discovery_service()
+        return self._proxy_location
+
     def _get_stub(self) -> PythonPanelServiceStub:
         if self._stub is None:
             if self._grpc_channel is not None:
                 self._stub = PythonPanelServiceStub(self._grpc_channel)
             else:
-                with self._initialization_lock:
-                    if self._grpc_channel_pool is None:
-                        _logger.debug("Creating unshared GrpcChannelPool.")
-                        self._grpc_channel_pool = GrpcChannelPool()
-                    if self._discovery_client is None:
-                        _logger.debug("Creating unshared DiscoveryClient.")
-                        self._discovery_client = DiscoveryClient(
-                            grpc_channel_pool=self._grpc_channel_pool
-                        )
-
-                    service_location = self._discovery_client.resolve_service(
-                        provided_interface=self._provided_interface,
-                        service_class=self._service_class,
-                    )
-                    channel = self._grpc_channel_pool.get_channel(service_location.insecure_address)
-                    self._stub = PythonPanelServiceStub(channel)
+                channel = self._get_panel_service_channel()
+                self._stub = PythonPanelServiceStub(channel)
         return self._stub
 
+    def _get_grpc_channel_pool(self) -> GrpcChannelPool:
+        if self._grpc_channel_pool is None:
+            _logger.debug("Creating unshared GrpcChannelPool.")
+            self._grpc_channel_pool = GrpcChannelPool()
+        return self._grpc_channel_pool
+
+    def _get_discovery_client(self, channel_pool: GrpcChannelPool) -> DiscoveryClient:
+        if self._discovery_client is None:
+            _logger.debug("Creating unshared DiscoveryClient.")
+            self._discovery_client = DiscoveryClient(grpc_channel_pool=channel_pool)
+        return self._discovery_client
+
+    def _resolve_service(
+        self, discovery_client: DiscoveryClient, provided_interface: str, service_class: str = ""
+    ) -> ServiceLocation:
+        _logger.debug("Resolving '%s'.", provided_interface)
+        service_location = discovery_client.resolve_service(
+            provided_interface=provided_interface,
+            service_class=service_class,
+        )
+        return service_location
+
+    def _get_panel_service_channel(self) -> grpc.Channel:
+        with self._initialization_lock:
+            grpc_channel_pool = self._get_grpc_channel_pool()
+            discovery_client = self._get_discovery_client(grpc_channel_pool)
+            panel_location = self._resolve_service(
+                discovery_client, self._provided_interface, self._service_class
+            )
+            channel = grpc_channel_pool.get_channel(panel_location.insecure_address)
+            return channel
+
+    def _resolve_discovery_service(self) -> ServiceLocation:
+        with self._initialization_lock:
+            grpc_channel_pool = self._get_grpc_channel_pool()
+            discovery_client = self._get_discovery_client(grpc_channel_pool)
+            discovery_location = self._resolve_service(
+                discovery_client, provided_interface="ni.http1.proxy"
+            )
+            return discovery_location
+
     def _invoke_with_retry(
         self, method: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
     ) -> _T:

src/nipanel/_panel_value_accessor.py

@@ -6,7 +6,7 @@
 from typing import TypeVar, overload
 
 import grpc
-from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
+from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient, ServiceLocation
 from ni_measurement_plugin_sdk_service.grpc.channelpool import GrpcChannelPool
 
 from nipanel._panel_client import PanelClient
@@ -55,6 +55,11 @@ def panel_id(self) -> str:
         """Read-only accessor for the panel ID."""
         return self._panel_id
 
+    @property
+    def proxy_location(self) -> ServiceLocation:
+        """Read-only accessor for the panel's proxy location."""
+        return self._panel_client.proxy_location
+
     @overload
     def get_value(self, value_id: str) -> object: ...
 

src/nipanel/_streamlit_constants.py

@@ -1,2 +1 @@
 STREAMLIT_PYTHON_PANEL_SERVICE = "ni.pythonpanel.v1.PythonPanelService"
-STREAMLIT_REFRESH_COMPONENT_URL = "http://localhost:42001/panel-service/refresh"

src/nipanel/_streamlit_panel_initializer.py

@@ -63,7 +63,8 @@ def get_panel_accessor() -> StreamlitPanelValueAccessor:
 
     panel = cast(StreamlitPanelValueAccessor, st.session_state[PANEL_ACCESSOR_KEY])
     _sync_session_state(panel)
-    refresh_component = initialize_refresh_component(panel.panel_id)
+    proxy_url = panel.proxy_location.insecure_address
+    refresh_component = initialize_refresh_component(proxy_url, panel.panel_id)
     refresh_component()
     return panel
 

src/nipanel/streamlit_refresh/__init__.py

@@ -1,17 +1,14 @@
 """Initializes a refresh component for Streamlit."""
 
-from typing import Any
+from streamlit.components.v1 import declare_component
+from streamlit.components.v1.custom_component import CustomComponent
 
-import streamlit.components.v1 as components
 
-from nipanel._streamlit_constants import STREAMLIT_REFRESH_COMPONENT_URL
-
-
-def initialize_refresh_component(panel_id: str) -> Any:
+def initialize_refresh_component(proxy_url: str, panel_id: str) -> CustomComponent:
     """Initialize a refresh component to the Streamlit app."""
-    _refresh_component_func = components.declare_component(
+    _refresh_component_func = declare_component(
         "panelRefreshComponent",
-        url=f"{STREAMLIT_REFRESH_COMPONENT_URL}/{panel_id}",
+        url=f"http://{proxy_url}/panel-service/refresh/{panel_id}",
     )
 
     return _refresh_component_func