diff --git a/src/datajoint/diagram.py b/src/datajoint/diagram.py index 20cd34b58..c51109edb 100644 --- a/src/datajoint/diagram.py +++ b/src/datajoint/diagram.py @@ -1496,8 +1496,19 @@ def make_dot(self): dot.set_rankdir(direction) for node in dot.get_nodes(): node.set_shape("circle") - name = node.get_name().strip('"') - props = node_props[name] + # Recover the original node key. `_encapsulate_node_names` wraps every + # key in one layer of double quotes; a plain strip('"') over-strips + # when the key is itself a quoted PostgreSQL table name + # (`"schema"."table"`), which previously raised a KeyError below + # (#1535). Try candidates and pick the one that is an actual key. + raw = node.get_name() + name = next( + (k for k in (raw, raw.strip('"'), raw[1:-1]) if k in node_props), + raw.strip('"'), + ) + props = node_props.get(name) + if props is None: + continue # unknown node — leave default styling rather than crash node.set_fontsize(props["fontsize"]) node.set_fontcolor(props["fontcolor"]) node.set_shape(props["shape"]) diff --git a/tests/integration/test_diagram_nocontext.py b/tests/integration/test_diagram_nocontext.py new file mode 100644 index 000000000..5809b8aed --- /dev/null +++ b/tests/integration/test_diagram_nocontext.py @@ -0,0 +1,54 @@ +""" +Regression test for #1535: dj.Diagram must render even when class-name +resolution fails and nodes fall back to raw (quoted) table names — previously a +KeyError on PostgreSQL, where `"schema"."table"` node keys were over-stripped. +""" + +import time + +import pytest + +import datajoint as dj + + +@pytest.fixture(scope="function") +def schema_by_backend(connection_by_backend, db_creds_by_backend): + backend = db_creds_by_backend["backend"] + test_id = str(int(time.time() * 1000))[-8:] + schema_name = f"djtest_noctx_{backend}_{test_id}"[:64] + if connection_by_backend.is_connected: + try: + connection_by_backend.query( + f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)}" + ) + except Exception: + pass + schema = dj.Schema(schema_name, connection=connection_by_backend) + yield schema + if connection_by_backend.is_connected: + try: + connection_by_backend.query( + f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)}" + ) + except Exception: + pass + + +def test_diagram_renders_without_context(schema_by_backend): + """With an empty context, nodes keep raw table names; rendering must still + succeed (no KeyError) on both backends.""" + if not dj.diagram.diagram_active: + pytest.skip("networkx/pydot not available") + + @schema_by_backend + class Master(dj.Manual): + definition = "master_id : int32" + + class Part(dj.Part): + definition = "-> master\npart_id : int32" + + # Empty context => class-name resolution fails => raw table-name nodes. + dot = dj.Diagram(schema_by_backend, context={}).make_dot() + names = [n.get_name() for n in dot.get_nodes()] + # The master's table appears as a node (by raw table name). + assert any("master" in n.lower() for n in names), f"master node missing; nodes={names}"