diff --git a/src/datajoint/heading.py b/src/datajoint/heading.py index 9addae474..0c126c533 100644 --- a/src/datajoint/heading.py +++ b/src/datajoint/heading.py @@ -765,6 +765,7 @@ def set_primary_key(self, primary_key): (dict(self.attributes[name].todict(), in_key=True) for name in primary_key), (dict(self.attributes[name].todict(), in_key=False) for name in self.names if name not in primary_key), ), + table_info=self.table_info, lineage_available=self._lineage_available, ) diff --git a/tests/unit/test_pk_rules.py b/tests/unit/test_pk_rules.py index 2c554091b..148bbacc1 100644 --- a/tests/unit/test_pk_rules.py +++ b/tests/unit/test_pk_rules.py @@ -5,6 +5,9 @@ primary keys in join operations. """ +from types import SimpleNamespace + +from datajoint.adapters import PostgreSQLAdapter from datajoint.heading import Heading @@ -205,3 +208,14 @@ def test_b_determines_a_ordering(self): assert names.index("x") < names.index("z") # y should be secondary (demoted from A's PK) assert "y" in result.secondary_attributes + + +class TestSetPrimaryKey: + """Tests for Heading.set_primary_key.""" + + def test_keeps_the_adapter_for_a_later_rename(self): + """A rename after set_primary_key (as dj.U & table does) quotes for the connection's backend.""" + heading = make_heading(["a"], ["b"]) + heading.table_info = {"conn": SimpleNamespace(adapter=PostgreSQLAdapter())} + renamed = heading.set_primary_key(["b"]).select([], rename_map={"c": "b"}) + assert renamed.attributes["c"].attribute_expression == '"b"'