Implement stub callback for node expansion

Author: ericfe-google

bigquery_magics/bigquery.py

@@ -596,21 +596,33 @@ def _handle_result(result, args):
     return result
 
 
-def _is_colab() -> bool:
-    """Check if code is running in Google Colab"""
-    try:
-        import google.colab  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
-def _colab_callback(query: str, params: str):
+def _colab_query_callback(query: str, params: str):
     return IPython.core.display.JSON(
         graph_server.convert_graph_data(query_results=json.loads(params))
     )
 
+def _colab_node_expansion_callback(request: dict, params_str: str):
+    """Handle node expansion requests in Google Colab environment
+    
+    Args:
+        request: A dictionary containing node expansion details including:
+            - uid: str - Unique identifier of the node to expand
+            - node_labels: List[str] - Labels of the node
+            - node_properties: List[Dict] - Properties of the node with key, value, and type
+            - direction: str - Direction of expansion ("INCOMING" or "OUTGOING")
+            - edge_label: Optional[str] - Label of edges to filter by
+        params_str: A JSON string containing connection parameters
+    
+    Returns:
+        JSON: A JSON-serialized response containing either:
+            - The query results with nodes and edges
+            - An error message if the request failed
+    """
+    try:
+        return IPython.core.display.JSON(graph_server.execute_node_expansion(params_str, request))
+    except BaseException as e:
+        return IPython.core.display.JSON({"error": e})
+
 
 singleton_server_thread: threading.Thread = None
 
@@ -628,11 +640,12 @@ def _add_graph_widget(query_result):
     # visualizer widget. In colab, we are not able to create an http server on a
     # background thread, so we use a special colab-specific api to register a callback,
     # to be invoked from Javascript.
-    if _is_colab():
+    try:
         from google.colab import output
 
-        output.register_callback("graph_visualization.Query", _colab_callback)
-    else:
+        output.register_callback("graph_visualization.Query", _colab_query_callback)
+        output.register_callback("graph_visualization.NodeExpansion", _colab_node_expansion_callback)
+    except ImportError:
         global singleton_server_thread
         alive = singleton_server_thread and singleton_server_thread.is_alive()
         if not alive:

bigquery_magics/graph_server.py

@@ -20,6 +20,10 @@
 from typing import Dict, List
 
 
+def execute_node_expansion(params, request):
+    return {"error": "Node expansion not yet implemented"}
+
+
 def convert_graph_data(query_results: Dict[str, Dict[str, str]]):
     """
     Converts graph data to the form expected by the visualization framework.
@@ -75,7 +79,7 @@ def convert_graph_data(query_results: Dict[str, Dict[str, str]]):
                     row_json = json.loads(value_value)
                     data[column_name].append(row_json)
                     tabular_data[column_name].append(row_json)
-                except:
+                except (ValueError, TypeError):
                     # Non-JSON columns cannot be visualized, but we still want them
                     # in the tabular view.
                     tabular_data[column_name].append(str(value_value))
@@ -113,6 +117,7 @@ class GraphServer:
     endpoints = {
         "get_ping": "/get_ping",
         "post_ping": "/post_ping",
+        "post_node_expansion": '/post_node_expansion',
         "post_query": "/post_query",
     }
 

tests/unit/test_bigquery.py

@@ -891,8 +891,8 @@ def test_bigquery_graph_colab(monkeypatch):
     graph_visualization is None or bigquery_storage is None,
     reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
 )
-def test_colab_callback():
-    result = bigquery_magics.bigquery._colab_callback(
+def test_colab_query_callback():
+    result = bigquery_magics.bigquery._colab_query_callback(
         "query", json.dumps({"result": {}})
     )
     assert result.data == {
@@ -905,6 +905,23 @@ def test_colab_callback():
     }
 
 
+@pytest.mark.usefixtures("ipython_interactive")
+@pytest.mark.skipif(
+    graph_visualization is None or bigquery_storage is None,
+    reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
+)
+def test_colab_node_expansion_callback():
+    result = bigquery_magics.bigquery._colab_node_expansion_callback(
+        request={"uid": "test_uid",
+                 "node_labels": ["label1, label2"],
+                 "node_properites": {},
+                 "direction": "INCOMING",
+                 "edge_label": None},
+        params_str="{}")
+
+    assert result.data == {"error": "Node expansion not yet implemented"}
+
+
 @pytest.mark.usefixtures("ipython_interactive")
 @pytest.mark.skipif(
     graph_visualization is not None or bigquery_storage is None,