Skip to content

Commit ed89a30

Browse files
fix(postgres): read primary and foreign keys from pg_catalog
PostgreSQL's `information_schema.table_constraints` lists only constraints on tables the current user owns or holds a privilege other than SELECT on. A login with only SELECT therefore saw no primary key on any table: `Table.primary_key` was empty, `parents(primary=True)` returned nothing so `key_source` raised, and `len()` built `count(DISTINCT )`. MySQL was unaffected. The keys are now read from `pg_constraint` and `pg_attribute`, which every login can read, the same way `load_foreign_keys_sql` already did: - `get_primary_key_sql` (the heading, `Table.primary_key`) and `load_primary_keys_sql` (the dependency graph) select `contype = 'p'`, ordered by position within the constraint -- the same column order `information_schema` reported. - `get_foreign_keys_sql` had the same problem, worse: `constraint_column_usage` shows only tables the user owns. It also joined the two column lists on constraint name alone, so a multi-column foreign key came back as a cross product of its columns. It now returns one row per foreign-key column, pairing each with the column it references, as the MySQL adapter does. Nothing in the library calls it today; output column names are unchanged. The new integration test creates a SELECT-only login on each backend and checks that primary keys, headings, parents and children, the dependency graph, `key_source`, and `get_foreign_keys_sql` all match what the owner sees. Against the previous code all five PostgreSQL cases fail; MySQL passed before and after. Closes #1554.
1 parent a00b89f commit ed89a30

2 files changed

Lines changed: 208 additions & 34 deletions

File tree

‎src/datajoint/adapters/postgres.py‎

Lines changed: 41 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -791,52 +791,59 @@ def get_columns_sql(self, schema_name: str, table_name: str) -> str:
791791
f"ORDER BY c.ordinal_position"
792792
)
793793

794+
# Key and constraint queries read the system catalogs (pg_constraint,
795+
# pg_attribute) rather than information_schema. information_schema hides
796+
# constraints on tables where the current user holds only SELECT, so a
797+
# read-only login would see no primary keys and no foreign keys. The
798+
# catalogs are readable by every login.
799+
794800
def get_primary_key_sql(self, schema_name: str, table_name: str) -> str:
795-
"""Query to get primary key columns."""
801+
"""Query to get primary key columns, in key order."""
796802
return (
797-
f"SELECT column_name FROM information_schema.key_column_usage "
798-
f"WHERE table_schema = {self.quote_string(schema_name)} "
799-
f"AND table_name = {self.quote_string(table_name)} "
800-
f"AND constraint_name IN ("
801-
f" SELECT constraint_name FROM information_schema.table_constraints "
802-
f" WHERE table_schema = {self.quote_string(schema_name)} "
803-
f" AND table_name = {self.quote_string(table_name)} "
804-
f" AND constraint_type = 'PRIMARY KEY'"
805-
f") "
806-
f"ORDER BY ordinal_position"
803+
f"SELECT a.attname AS column_name "
804+
f"FROM pg_constraint c "
805+
f"JOIN pg_class cl ON c.conrelid = cl.oid "
806+
f"JOIN pg_namespace ns ON cl.relnamespace = ns.oid "
807+
f"CROSS JOIN LATERAL unnest(c.conkey) WITH ORDINALITY AS cols(conkey, ord) "
808+
f"JOIN pg_attribute a ON a.attrelid = cl.oid AND a.attnum = cols.conkey "
809+
f"WHERE c.contype = 'p' "
810+
f"AND ns.nspname = {self.quote_string(schema_name)} "
811+
f"AND cl.relname = {self.quote_string(table_name)} "
812+
f"ORDER BY cols.ord"
807813
)
808814

809815
def get_foreign_keys_sql(self, schema_name: str, table_name: str) -> str:
810816
"""Query to get foreign key constraints."""
811817
return (
812-
f"SELECT kcu.constraint_name, kcu.column_name, "
813-
f"ccu.table_name AS foreign_table_name, ccu.column_name AS foreign_column_name "
814-
f"FROM information_schema.key_column_usage AS kcu "
815-
f"JOIN information_schema.constraint_column_usage AS ccu "
816-
f" ON kcu.constraint_name = ccu.constraint_name "
817-
f"WHERE kcu.table_schema = {self.quote_string(schema_name)} "
818-
f"AND kcu.table_name = {self.quote_string(table_name)} "
819-
f"AND kcu.constraint_name IN ("
820-
f" SELECT constraint_name FROM information_schema.table_constraints "
821-
f" WHERE table_schema = {self.quote_string(schema_name)} "
822-
f" AND table_name = {self.quote_string(table_name)} "
823-
f" AND constraint_type = 'FOREIGN KEY'"
824-
f") "
825-
f"ORDER BY kcu.constraint_name, kcu.ordinal_position"
818+
f"SELECT c.conname AS constraint_name, a1.attname AS column_name, "
819+
f"cl2.relname AS foreign_table_name, a2.attname AS foreign_column_name "
820+
f"FROM pg_constraint c "
821+
f"JOIN pg_class cl1 ON c.conrelid = cl1.oid "
822+
f"JOIN pg_namespace ns1 ON cl1.relnamespace = ns1.oid "
823+
f"JOIN pg_class cl2 ON c.confrelid = cl2.oid "
824+
f"CROSS JOIN LATERAL unnest(c.conkey, c.confkey) WITH ORDINALITY AS cols(conkey, confkey, ord) "
825+
f"JOIN pg_attribute a1 ON a1.attrelid = cl1.oid AND a1.attnum = cols.conkey "
826+
f"JOIN pg_attribute a2 ON a2.attrelid = cl2.oid AND a2.attnum = cols.confkey "
827+
f"WHERE c.contype = 'f' "
828+
f"AND ns1.nspname = {self.quote_string(schema_name)} "
829+
f"AND cl1.relname = {self.quote_string(table_name)} "
830+
f"ORDER BY c.conname, cols.ord"
826831
)
827832

828833
def load_primary_keys_sql(self, schemas_list: str, like_pattern: str) -> str:
829834
"""Query to load all primary key columns across schemas."""
830-
tab_expr = "'\"' || kcu.table_schema || '\".\"' || kcu.table_name || '\"'"
835+
tab_expr = "'\"' || ns.nspname || '\".\"' || cl.relname || '\"'"
831836
return (
832-
f"SELECT {tab_expr} as tab, kcu.column_name "
833-
f"FROM information_schema.key_column_usage kcu "
834-
f"JOIN information_schema.table_constraints tc "
835-
f"ON kcu.constraint_name = tc.constraint_name "
836-
f"AND kcu.table_schema = tc.table_schema "
837-
f"WHERE kcu.table_name NOT LIKE {like_pattern} "
838-
f"AND kcu.table_schema in ({schemas_list}) "
839-
f"AND tc.constraint_type = 'PRIMARY KEY'"
837+
f"SELECT {tab_expr} as tab, a.attname as column_name "
838+
f"FROM pg_constraint c "
839+
f"JOIN pg_class cl ON c.conrelid = cl.oid "
840+
f"JOIN pg_namespace ns ON cl.relnamespace = ns.oid "
841+
f"CROSS JOIN LATERAL unnest(c.conkey) WITH ORDINALITY AS cols(conkey, ord) "
842+
f"JOIN pg_attribute a ON a.attrelid = cl.oid AND a.attnum = cols.conkey "
843+
f"WHERE c.contype = 'p' "
844+
f"AND cl.relname NOT LIKE {like_pattern} "
845+
f"AND ns.nspname in ({schemas_list}) "
846+
f"ORDER BY ns.nspname, cl.relname, cols.ord"
840847
)
841848

842849
def load_foreign_keys_sql(self, schemas_list: str, like_pattern: str) -> str:
Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
"""
2+
A login holding only SELECT on a schema must see the same table structure as
3+
the schema's owner: primary keys, headings, and the dependency graph (#1554).
4+
5+
On PostgreSQL, information_schema hides constraints on tables where the current
6+
user holds only SELECT, so key and foreign-key metadata is read from the system
7+
catalogs instead. The same checks run on MySQL, where a SELECT-only login has
8+
always seen this metadata.
9+
"""
10+
11+
import pytest
12+
13+
import datajoint as dj
14+
15+
READER = "djreadonly"
16+
TABLES = ("Rig", "Subject", "Session", "Pairing", "Analysis")
17+
18+
19+
@pytest.fixture
20+
def owner_schema(connection_by_backend, backend, prefix):
21+
"""A small pipeline declared and populated by the owner."""
22+
schema = dj.Schema(f"{prefix}_read_only_{backend}", connection=connection_by_backend)
23+
24+
@schema
25+
class Rig(dj.Lookup):
26+
definition = """
27+
rig : varchar(16)
28+
"""
29+
contents = [("A",), ("B",)]
30+
31+
@schema
32+
class Subject(dj.Manual):
33+
definition = """
34+
subject_id : int32
35+
---
36+
species : varchar(32)
37+
"""
38+
39+
@schema
40+
class Session(dj.Manual):
41+
# Primary key order is not alphabetical, so a reordered key would show.
42+
definition = """
43+
-> Subject
44+
session_idx : int16
45+
---
46+
-> Rig
47+
note = '' : varchar(255)
48+
"""
49+
50+
@schema
51+
class Pairing(dj.Manual):
52+
definition = """
53+
-> Subject
54+
-> Subject.proj(partner_id='subject_id')
55+
"""
56+
57+
@schema
58+
class Analysis(dj.Computed):
59+
definition = """
60+
-> Session
61+
---
62+
result : float64
63+
"""
64+
65+
def make(self, key):
66+
self.insert1(dict(key, result=1.0))
67+
68+
Subject.insert([(1, "mouse"), (2, "rat")], skip_duplicates=True)
69+
Session.insert([(1, 1, "A", ""), (1, 2, "B", ""), (2, 1, "A", "")], skip_duplicates=True)
70+
Pairing.insert([(1, 2)], skip_duplicates=True)
71+
Analysis.populate(Session & "session_idx = 1")
72+
73+
yield schema
74+
schema.drop()
75+
76+
77+
def _drop_reader(conn, backend):
78+
if backend == "postgresql":
79+
if conn.query("SELECT 1 FROM pg_roles WHERE rolname = %s", args=(READER,)).fetchone():
80+
conn.query(f"DROP OWNED BY {READER}")
81+
conn.query(f"DROP ROLE {READER}")
82+
else:
83+
conn.query(f"DROP USER IF EXISTS '{READER}'@'%%'")
84+
85+
86+
@pytest.fixture
87+
def reader_connection(owner_schema, connection_by_backend, db_creds_by_backend, backend):
88+
"""A connection whose login holds only SELECT on the owner's schema."""
89+
owner = connection_by_backend
90+
database = owner.adapter.quote_identifier(owner_schema.database)
91+
_drop_reader(owner, backend)
92+
if backend == "postgresql":
93+
owner.query(f"CREATE ROLE {READER} LOGIN PASSWORD '{READER}'")
94+
owner.query(f"GRANT USAGE ON SCHEMA {database} TO {READER}")
95+
owner.query(f"GRANT SELECT ON ALL TABLES IN SCHEMA {database} TO {READER}")
96+
else:
97+
owner.query(f"CREATE USER '{READER}'@'%%' IDENTIFIED BY '{READER}'")
98+
owner.query(f"GRANT SELECT ON {database}.* TO '{READER}'@'%%'")
99+
100+
reader = dj.Connection(host=db_creds_by_backend["host"], user=READER, password=READER)
101+
yield reader
102+
reader.close()
103+
_drop_reader(owner, backend)
104+
105+
106+
@pytest.fixture
107+
def owner_and_reader(owner_schema, connection_by_backend, reader_connection):
108+
"""The same schema loaded as virtual modules by the owner and by the reader."""
109+
owner = dj.VirtualModule("owner", owner_schema.database, connection=connection_by_backend)
110+
reader = dj.VirtualModule("reader", owner_schema.database, connection=reader_connection)
111+
return owner, reader
112+
113+
114+
def test_primary_keys(owner_and_reader):
115+
owner, reader = owner_and_reader
116+
for name in TABLES:
117+
assert getattr(owner, name).primary_key, name
118+
assert getattr(reader, name).primary_key == getattr(owner, name).primary_key, name
119+
assert reader.Session.primary_key == ["subject_id", "session_idx"]
120+
121+
122+
def test_headings(owner_and_reader):
123+
owner, reader = owner_and_reader
124+
for name in TABLES:
125+
assert getattr(reader, name).heading.attributes == getattr(owner, name).heading.attributes, name
126+
127+
128+
def _graph(conn):
129+
conn.dependencies.load()
130+
nodes = dict(conn.dependencies.nodes(data=True))
131+
edges = {(u, v, k): d for u, v, k, d in conn.dependencies.edges(keys=True, data=True)}
132+
return nodes, edges
133+
134+
135+
def test_dependencies(owner_and_reader, connection_by_backend, reader_connection):
136+
owner, reader = owner_and_reader
137+
for name in TABLES:
138+
owner_table, reader_table = getattr(owner, name), getattr(reader, name)
139+
for primary in (None, True, False):
140+
assert reader_table.parents(primary=primary) == owner_table.parents(primary=primary), (name, primary)
141+
assert reader_table.children(primary=primary) == owner_table.children(primary=primary), (name, primary)
142+
143+
nodes, edges = _graph(connection_by_backend)
144+
assert _graph(reader_connection) == (nodes, edges)
145+
assert len(edges) == 5 # Subject->Session, Rig->Session, Subject->Pairing (x2), Session->Analysis
146+
147+
148+
def test_queries_that_need_keys(owner_and_reader):
149+
owner, reader = owner_and_reader
150+
assert len(reader.Session()) == len(owner.Session()) == 3
151+
assert len(reader.Analysis.key_source) == 3
152+
assert len(reader.Analysis.key_source - reader.Analysis) == 1
153+
assert reader.Session.keys(order_by="KEY") == owner.Session.keys(order_by="KEY")
154+
assert reader.Session.describe() == owner.Session.describe()
155+
156+
157+
def test_foreign_keys_sql(owner_and_reader, connection_by_backend, reader_connection):
158+
"""``get_foreign_keys_sql`` returns one row per foreign-key column, for any login."""
159+
owner, _ = owner_and_reader
160+
adapter = connection_by_backend.adapter
161+
rows = {}
162+
for name in TABLES:
163+
table = getattr(owner, name)
164+
sql = adapter.get_foreign_keys_sql(table.database, table.table_name)
165+
rows[name] = connection_by_backend.query(sql).fetchall()
166+
assert reader_connection.query(sql).fetchall() == rows[name], name
167+
assert [len(rows[name]) for name in TABLES] == [0, 0, 2, 2, 2]

0 commit comments

Comments
 (0)