Skip to content

Commit 54e3419

Browse files
fix(#1535): render diagrams with unresolved (raw table-name) nodes
dj.Diagram(...).make_dot() raised a KeyError on PostgreSQL whenever a node's name could not be resolved to a class name and fell back to the raw quoted table name ("schema"."table"): _encapsulate_node_names adds one layer of quotes, but the styling loop's strip('"') removed all of them, over-stripping the pg key so node_props[name] missed. Recover the real key by trying candidates (raw / all- stripped / one-layer-stripped) and picking the one that is an actual key; degrade to default styling instead of crashing if none match. MySQL was unaffected (backtick quoting survives strip). Adds a both-backends regression test.
1 parent 8d05d98 commit 54e3419

2 files changed

Lines changed: 67 additions & 2 deletions

File tree

‎src/datajoint/diagram.py‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1496,8 +1496,19 @@ def make_dot(self):
14961496
dot.set_rankdir(direction)
14971497
for node in dot.get_nodes():
14981498
node.set_shape("circle")
1499-
name = node.get_name().strip('"')
1500-
props = node_props[name]
1499+
# Recover the original node key. `_encapsulate_node_names` wraps every
1500+
# key in one layer of double quotes; a plain strip('"') over-strips
1501+
# when the key is itself a quoted PostgreSQL table name
1502+
# (`"schema"."table"`), which previously raised a KeyError below
1503+
# (#1535). Try candidates and pick the one that is an actual key.
1504+
raw = node.get_name()
1505+
name = next(
1506+
(k for k in (raw, raw.strip('"'), raw[1:-1]) if k in node_props),
1507+
raw.strip('"'),
1508+
)
1509+
props = node_props.get(name)
1510+
if props is None:
1511+
continue # unknown node — leave default styling rather than crash
15011512
node.set_fontsize(props["fontsize"])
15021513
node.set_fontcolor(props["fontcolor"])
15031514
node.set_shape(props["shape"])
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
"""
2+
Regression test for #1535: dj.Diagram must render even when class-name
3+
resolution fails and nodes fall back to raw (quoted) table names — previously a
4+
KeyError on PostgreSQL, where `"schema"."table"` node keys were over-stripped.
5+
"""
6+
7+
import time
8+
9+
import pytest
10+
11+
import datajoint as dj
12+
13+
14+
@pytest.fixture(scope="function")
15+
def schema_by_backend(connection_by_backend, db_creds_by_backend):
16+
backend = db_creds_by_backend["backend"]
17+
test_id = str(int(time.time() * 1000))[-8:]
18+
schema_name = f"djtest_noctx_{backend}_{test_id}"[:64]
19+
if connection_by_backend.is_connected:
20+
try:
21+
connection_by_backend.query(
22+
f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)}"
23+
)
24+
except Exception:
25+
pass
26+
schema = dj.Schema(schema_name, connection=connection_by_backend)
27+
yield schema
28+
if connection_by_backend.is_connected:
29+
try:
30+
connection_by_backend.query(
31+
f"DROP DATABASE IF EXISTS {connection_by_backend.adapter.quote_identifier(schema_name)}"
32+
)
33+
except Exception:
34+
pass
35+
36+
37+
def test_diagram_renders_without_context(schema_by_backend):
38+
"""With an empty context, nodes keep raw table names; rendering must still
39+
succeed (no KeyError) on both backends."""
40+
if not dj.diagram.diagram_active:
41+
pytest.skip("networkx/pydot not available")
42+
43+
@schema_by_backend
44+
class Master(dj.Manual):
45+
definition = "master_id : int32"
46+
47+
class Part(dj.Part):
48+
definition = "-> master\npart_id : int32"
49+
50+
# Empty context => class-name resolution fails => raw table-name nodes.
51+
dot = dj.Diagram(schema_by_backend, context={}).make_dot()
52+
names = [n.get_name() for n in dot.get_nodes()]
53+
# The master's table appears as a node (by raw table name).
54+
assert any("master" in n.lower() for n in names), f"master node missing; nodes={names}"

0 commit comments

Comments
 (0)