Add support for multiple ion servers in UI

CHANGES.md

@@ -1,5 +1,9 @@
 # Change Log
 
+### v0.15.0 - 2023-12-14
+
+* Added support for multiple Cesium ion servers by creating `CesiumIonServerPrim` prims.
+
 ### v0.14.0 - 2023-12-01
 
 * Added support for `EXT_structural_metadata`. Property values can be accessed in material graph with the `cesium_property` nodes.

exts/cesium.omniverse/cesium/omniverse/bindings/CesiumOmniversePythonBindings.pyi

@@ -73,12 +73,15 @@ class ICesiumOmniverseInterface:
     def create_token(self, arg0: str) -> None: ...
     def credits_available(self) -> bool: ...
     def credits_start_next_frame(self) -> None: ...
+    def get_all_server_paths(self) -> List[str]: ...
+    def get_all_sessions(self, *args, **kwargs) -> Any: ...
     def get_all_tileset_paths(self) -> List[str]: ...
     def get_asset_token_troubleshooting_details(self, *args, **kwargs) -> Any: ...
     def get_asset_troubleshooting_details(self, *args, **kwargs) -> Any: ...
     def get_credits(self) -> List[Tuple[str, bool]]: ...
     def get_default_token_troubleshooting_details(self, *args, **kwargs) -> Any: ...
     def get_render_statistics(self, *args, **kwargs) -> Any: ...
+    def get_server_path(self) -> str: ...
     def get_session(self, *args, **kwargs) -> Any: ...
     def get_set_default_token_result(self, *args, **kwargs) -> Any: ...
     def is_default_token_set(self) -> bool: ...

exts/cesium.omniverse/cesium/omniverse/extension.py

@@ -460,6 +460,7 @@ def _setup_ion_server_prims(self):
             # If we have no ion server prims, lets add a default one for the official ion servers.
             path = "/CesiumServers/IonOfficial"
             prim: CesiumIonServer = CesiumIonServer.Define(stage, path)
+            prim.GetDisplayNameAttr().Set("ion.cesium.com")
             prim.GetIonServerUrlAttr().Set("https://ion.cesium.com/")
             prim.GetIonServerApiUrlAttr().Set("https://api.cesium.com/")
             prim.GetIonServerApplicationIdAttr().Set(413)

exts/cesium.omniverse/cesium/omniverse/ui/attributes/ion_server_attributes_widget.py

@@ -21,6 +21,7 @@ def _customize_props_layout(self, props):
 
         with frame:
             with CustomLayoutGroup("ion Server"):
+                CustomLayoutProperty("cesium:displayName")
                 CustomLayoutProperty("cesium:ionServerUrl")
                 CustomLayoutProperty("cesium:ionServerApiUrl")
                 CustomLayoutProperty("cesium:ionServerApplicationId")

exts/cesium.omniverse/cesium/omniverse/ui/main_window.py

@@ -190,6 +190,8 @@ def _build_fn(self):
         with ui.VStack(spacing=0):
             button_style = CesiumOmniverseUiStyles.top_bar_button_style
 
+            self._profile_widget = CesiumOmniverseProfileWidget(self._cesium_omniverse_interface, height=20)
+
             with ui.HStack(height=ui.Length(80, ui.UnitType.PIXEL)):
                 self._add_button = ui.Button(
                     "Add",
@@ -233,8 +235,6 @@ def _build_fn(self):
                 )
             self._quick_add_widget = CesiumOmniverseQuickAddWidget(self._cesium_omniverse_interface)
             self._sign_in_widget = CesiumOmniverseSignInWidget(self._cesium_omniverse_interface, visible=False)
-            ui.Spacer()
-            self._profile_widget = CesiumOmniverseProfileWidget(self._cesium_omniverse_interface, height=20)
 
     def _add_button_clicked(self) -> None:
         if not self._add_button or not self._add_button.enabled:

exts/cesium.omniverse/cesium/omniverse/ui/profile_widget.py

@@ -1,23 +1,116 @@
+import omni.usd
 import logging
 import carb.events
 import omni.kit.app as app
 import omni.ui as ui
-import webbrowser
 from typing import List, Optional
-from ..bindings import ICesiumOmniverseInterface
+from ..bindings import ICesiumOmniverseInterface, CesiumIonSession
+from enum import Enum
+from cesium.usd.plugins.CesiumUsdSchemas import IonServer as CesiumIonServer
+from ..usdUtils import set_path_to_current_ion_server
 
-LOADING_PROFILE_MESSAGE = "Loading user information..."
-CONNECTED_TO_MESSAGE_BASE = "Connected to Cesium ion as"
+
+class SessionState(Enum):
+    NOT_CONNECTED = 1
+    LOADING = 2
+    CONNECTED = 3
+
+
+def get_session_state(session: CesiumIonSession) -> SessionState:
+    if session.is_profile_loaded():
+        return SessionState.CONNECTED
+    elif session.is_loading_profile():
+        return SessionState.LOADING
+    else:
+        return SessionState.NOT_CONNECTED
+
+
+def get_profile_id(session: CesiumIonSession) -> Optional[int]:
+    if session.is_profile_loaded():
+        profile = session.get_profile()
+        return profile.id
+
+    return None
+
+
+class SessionComboItem(ui.AbstractItem):
+    def __init__(self, session: CesiumIonSession, server: CesiumIonServer):
+        super().__init__()
+
+        session_state = get_session_state(session)
+        prefix = ""
+        suffix = ""
+
+        if session_state == SessionState.NOT_CONNECTED:
+            suffix += " (not connected)"
+        elif session_state == SessionState.LOADING:
+            suffix += " (loading profile...)"
+        elif session_state == SessionState.CONNECTED:
+            prefix += session.get_profile().username
+            prefix += " @ "
+
+        # Get the display name from the server prim. If that's empty, use the prim path.
+        server_name = server.GetDisplayNameAttr().Get()
+        if server_name == "":
+            server_name = server.GetPath()
+
+        self.text = ui.SimpleStringModel(f"{prefix}{server_name}{suffix}")
+        self.server = server
+
+
+class SessionComboModel(ui.AbstractItemModel):
+    def __init__(self):
+        super().__init__()
+        self._logger = logging.getLogger(__name__)
+
+        self._current_index = ui.SimpleIntModel(0)
+        self._current_index.add_value_changed_fn(lambda index_model: self._item_changed(None))
+
+        self._items = []
+
+    def replace_all_items(
+        self, sessions: List[CesiumIonSession], servers: List[CesiumIonServer], current_server: CesiumIonServer
+    ):
+        self._items.clear()
+        self._items = [SessionComboItem(session, server) for session, server in zip(sessions, servers)]
+
+        current_index = 0
+        for index, server in enumerate(servers):
+            if server.GetPath() == current_server.GetPath():
+                current_index = index
+                break
+
+        self._current_index.set_value(current_index)
+        self._item_changed(None)
+
+    def get_item_children(self, item=None):
+        return self._items
+
+    def get_item_value_model(self, item: SessionComboItem = None, column_id: int = 0):
+        if item is None:
+            return self._current_index
+        return item.text
+
+    def get_current_selection(self):
+        if len(self._items) < 1:
+            return None
+
+        return self._items[self._current_index.get_value_as_int()]
 
 
 class CesiumOmniverseProfileWidget(ui.Frame):
     def __init__(self, cesium_omniverse_interface: ICesiumOmniverseInterface, **kwargs):
         self._logger = logging.getLogger(__name__)
         self._cesium_omniverse_interface = cesium_omniverse_interface
 
-        self._profile_id: Optional[int] = None
-        self._button_enabled = False
-        self._message = ""
+        self._profile_ids: List[int] = []
+        self._session_states: List[SessionState] = []
+        self._server_paths: List[str] = []
+        self._server_names: List[str] = []
+
+        self._sessions_combo_box: Optional[ui.ComboBox] = None
+        self._sessions_combo_model = SessionComboModel()
+        self._sessions_combo_model.add_item_changed_fn(self._on_item_changed)
 
         self._subscriptions: List[carb.events.ISubscription] = []
         self._setup_subscriptions()
@@ -35,33 +128,54 @@ def _setup_subscriptions(self):
             update_stream.create_subscription_to_pop(self._on_update_frame, name="on_update_frame")
         )
 
+    def _on_item_changed(self, item_model, item):
+        item = self._sessions_combo_model.get_current_selection()
+        if item is not None:
+            server_path = item.server.GetPath()
+            set_path_to_current_ion_server(server_path)
+
     def _on_update_frame(self, _e: carb.events.IEvent):
-        session = self._cesium_omniverse_interface.get_session()
-        if session is not None:
-            if session.is_profile_loaded():
-                profile = session.get_profile()
-                if self._profile_id != profile.id:
-                    self._profile_id = profile.id
-                    self._button_enabled = True
-                    self._message = f"{CONNECTED_TO_MESSAGE_BASE} {profile.username}"
-                    self.rebuild()
-            elif session.is_loading_profile():
-                self.visible = True
-                self._button_enabled = False
-                self._message = LOADING_PROFILE_MESSAGE
-                self.rebuild()
-            elif session.is_connected():
+        if omni.usd.get_context().get_stage_state() != omni.usd.StageState.OPENED:
+            return
+
+        stage = omni.usd.get_context().get_stage()
+
+        sessions = self._cesium_omniverse_interface.get_all_sessions()
+        server_paths = self._cesium_omniverse_interface.get_all_server_paths()
+        servers = [CesiumIonServer.Get(stage, server_path) for server_path in server_paths]
+        server_names = [server.GetDisplayNameAttr().Get() for server in servers]
+        current_server_path = self._cesium_omniverse_interface.get_server_path()
+        current_server = CesiumIonServer.Get(stage, current_server_path)
+
+        profile_ids = []
+        session_states = []
+
+        for session in sessions:
+            profile_id = get_profile_id(session)
+            session_state = get_session_state(session)
+
+            if session.is_connected() and not session.is_profile_loaded():
                 session.refresh_profile()
-            else:
-                self._profile_id = None
-                self.visible = False
 
-    def _on_profile_button_clicked(self) -> None:
-        if self.visible:
-            # We just open a link to ion directly.
-            # They may have to sign in if they aren't signed in on their browser already.
-            webbrowser.open("https://cesium.com/ion")
+            profile_ids.append(profile_id)
+            session_states.append(session_state)
+
+        if (
+            profile_ids != self._profile_ids
+            or session_states != self._session_states
+            or server_paths != self._server_paths
+            or server_names != self._server_names
+        ):
+            self._logger.info("Rebuilding profile widget")
+
+            self._profile_ids = profile_ids
+            self._session_states = session_states
+            self._server_paths = server_paths
+            self._server_names = server_names
+
+            self._sessions_combo_model.replace_all_items(sessions, servers, current_server)
+
+            self.rebuild()
 
     def _build_ui(self):
-        with self:
-            ui.Button(self._message, clicked_fn=self._on_profile_button_clicked, enabled=self._button_enabled)
+        self._sessions_combo_box = ui.ComboBox(self._sessions_combo_model)

exts/cesium.omniverse/cesium/omniverse/ui/quick_add_widget.py

@@ -2,10 +2,12 @@
 import carb.events
 import omni.kit.app as app
 import omni.ui as ui
+import omni.usd
 from typing import List, Optional
 from ..bindings import ICesiumOmniverseInterface
 from ..models import AssetToAdd
 from .styles import CesiumOmniverseUiStyles
+from cesium.usd.plugins.CesiumUsdSchemas import IonServer as CesiumIonServer
 
 LABEL_HEIGHT = 24
 BUTTON_HEIGHT = 40
@@ -38,10 +40,22 @@ def _on_update_frame(self, _: carb.events.IEvent):
         if self._ion_quick_add_frame is None:
             return
 
+        if omni.usd.get_context().get_stage_state() != omni.usd.StageState.OPENED:
+            return
+
         session = self._cesium_omniverse_interface.get_session()
 
         if session is not None:
-            self._ion_quick_add_frame.visible = session.is_connected()
+            stage = omni.usd.get_context().get_stage()
+            current_server_path = self._cesium_omniverse_interface.get_server_path()
+            current_server = CesiumIonServer.Get(stage, current_server_path)
+            current_server_url = current_server.GetIonServerUrlAttr().Get()
+
+            # Temporary workaround to only show quick add assets for official ion server
+            # until quick add route is implemented
+            self._ion_quick_add_frame.visible = (
+                session.is_connected() and current_server_url == "https://ion.cesium.com/"
+            )
 
     @staticmethod
     def _add_blank_button_clicked():

exts/cesium.omniverse/cesium/omniverse/usdUtils/usdUtils.py

@@ -1,14 +1,12 @@
 import omni.usd
 import omni.kit
 import re
-from typing import Optional
 
 from cesium.usd.plugins.CesiumUsdSchemas import (
     Imagery as CesiumImagery,
     Tileset as CesiumTileset,
     Tokens as CesiumTokens,
     Data as CesiumData,
-    IonServer as CesiumIonServer,
 )
 from pxr import Sdf
 from pxr.UsdGeom import Gprim
@@ -87,13 +85,13 @@ def get_path_to_current_ion_server() -> Sdf.Path:
     return targets[0]
 
 
-def get_current_ion_server_prim() -> Optional[CesiumIonServer]:
+def set_path_to_current_ion_server(server_path: str):
     stage = omni.usd.get_context().get_stage()
 
-    path = get_current_ion_server_prim()
+    data = CesiumData.Get(stage, CESIUM_DATA_PRIM_PATH)
 
-    prim: CesiumIonServer = CesiumIonServer.Get(stage, path)
-    if not prim.GetPrim().IsValid():
-        return None
+    if not data.GetPrim().IsValid():
+        return
 
-    return prim
+    rel = data.GetSelectedIonServerRel()
+    rel.SetTargets([server_path])

exts/cesium.usd.plugins/cesium/usd/plugins/CesiumUsdSchemas/__init__.pyi

@@ -193,6 +193,8 @@ class IonServer(pxr.Usd.Typed):
     @classmethod
     def __init__(cls, *args, **kwargs) -> None: ...
     @classmethod
+    def CreateDisplayNameAttr(cls, *args, **kwargs) -> Any: ...
+    @classmethod
     def CreateIonServerApiUrlAttr(cls, *args, **kwargs) -> Any: ...
     @classmethod
     def CreateIonServerApplicationIdAttr(cls, *args, **kwargs) -> Any: ...
@@ -207,6 +209,8 @@ class IonServer(pxr.Usd.Typed):
     @classmethod
     def Get(cls, *args, **kwargs) -> Any: ...
     @classmethod
+    def GetDisplayNameAttr(cls, *args, **kwargs) -> Any: ...
+    @classmethod
     def GetIonServerApiUrlAttr(cls, *args, **kwargs) -> Any: ...
     @classmethod
     def GetIonServerApplicationIdAttr(cls, *args, **kwargs) -> Any: ...
@@ -391,6 +395,8 @@ class Tokens(Boost.Python.instance):
     @property
     def cesiumDebugTexturePoolInitialCapacity(self) -> Any: ...
     @property
+    def cesiumDisplayName(self) -> Any: ...
+    @property
     def cesiumEcefToUsdTransform(self) -> Any: ...
     @property
     def cesiumEnableFogCulling(self) -> Any: ...