Skip to content

Commit db5f9e9

Browse files
marcelofernmarcelofern
committed
User must have REFERENCES privilege on referred tables
Prior to this change, there was no check to verify that the user had the REFERENCES privilege on the referred tables (tables the table to be repacked has a fk to). This is a problem because Psycopack needs to create new foreign keys on the shadow table that point to the referred tables, and that is not possible unless the user has REFERENCES privilege. This change adds an explicit check to ensure the user has this privilege for all referred tables before proceeding.
1 parent 62f2ab9 commit db5f9e9

4 files changed

Lines changed: 148 additions & 0 deletions

File tree

‎src/psycopack/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
InvalidPrimaryKeyTypeForConversion,
1212
InvalidStageForReset,
1313
NoCreateAndUsagePrivilegeOnSchema,
14+
NoReferencesPrivilege,
1415
NoReferringTableOwnership,
1516
NotTableOwner,
1617
PrimaryKeyNotFound,
@@ -34,6 +35,7 @@
3435
"InvalidPrimaryKeyTypeForConversion",
3536
"InvalidStageForReset",
3637
"NoCreateAndUsagePrivilegeOnSchema",
38+
"NoReferencesPrivilege",
3739
"NoReferringTableOwnership",
3840
"NotTableOwner",
3941
"PrimaryKeyNotFound",

‎src/psycopack/_introspect.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,14 @@ class ReferringForeignKey:
3333
is_owner: bool
3434

3535

36+
@dataclasses.dataclass
37+
class ReferredForeignKey:
38+
name: str
39+
schema: str
40+
privileges: list[str]
41+
is_owner: bool
42+
43+
3644
@dataclasses.dataclass
3745
class Trigger:
3846
name: str
@@ -290,6 +298,81 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
290298
for name, definition, is_validated, referring_table, schema, is_owner in results
291299
]
292300

301+
def get_referred_fks(self, *, table: str, schema: str) -> list[ReferredForeignKey]:
302+
self.cur.execute(
303+
psycopg.sql.SQL(
304+
dedent("""
305+
SELECT
306+
referred_pg_class.relname AS referred_table,
307+
referred_pg_namespace.nspname AS referred_schema,
308+
array_agg(
309+
table_privileges.privilege_type
310+
ORDER BY table_privileges.privilege_type
311+
) AS privileges,
312+
(
313+
SELECT
314+
tableowner = current_user
315+
FROM
316+
pg_tables
317+
WHERE
318+
schemaname = referred_pg_namespace.nspname
319+
AND tablename = referred_pg_class.relname
320+
) AS is_owner
321+
FROM
322+
pg_catalog.pg_constraint
323+
INNER JOIN
324+
pg_catalog.pg_class AS referring_pg_class
325+
ON (pg_constraint.conrelid = referring_pg_class.oid)
326+
INNER JOIN
327+
pg_catalog.pg_class AS referred_pg_class
328+
ON (pg_constraint.confrelid = referred_pg_class.oid)
329+
INNER JOIN
330+
pg_catalog.pg_namespace AS referring_pg_namespace
331+
ON (referring_pg_namespace.oid = referring_pg_class.relnamespace)
332+
INNER JOIN
333+
pg_catalog.pg_namespace AS referred_pg_namespace
334+
ON (referred_pg_namespace.oid = referred_pg_class.relnamespace)
335+
LEFT OUTER JOIN
336+
information_schema.table_privileges
337+
ON (
338+
table_privileges.table_name = referred_pg_class.relname
339+
AND table_privileges.table_schema = referred_pg_namespace.nspname
340+
)
341+
WHERE
342+
pg_constraint.conrelid = referring_pg_class.oid
343+
AND referring_pg_class.relname = {table}
344+
AND referring_pg_namespace.nspname = {schema}
345+
AND pg_constraint.contype = 'f'
346+
GROUP BY
347+
referred_pg_namespace.nspname,
348+
referred_pg_class.relname
349+
ORDER BY
350+
referred_pg_namespace.nspname,
351+
referred_pg_class.relname;
352+
""")
353+
)
354+
.format(
355+
table=psycopg.sql.Literal(table),
356+
schema=psycopg.sql.Literal(schema),
357+
)
358+
.as_string(self.conn)
359+
)
360+
results = self.cur.fetchall()
361+
assert results is not None
362+
return [
363+
ReferredForeignKey(
364+
name=name,
365+
schema=schema,
366+
privileges=(
367+
privileges.strip("{}").split(",")
368+
if "NULL" not in privileges
369+
else []
370+
),
371+
is_owner=is_owner,
372+
)
373+
for name, schema, privileges, is_owner in results
374+
]
375+
293376
def table_is_empty(self, *, table: str) -> int:
294377
self.cur.execute(
295378
psycopg.sql.SQL(

‎src/psycopack/_repack.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,10 @@ class NoReferringTableOwnership(BaseRepackError):
7676
pass
7777

7878

79+
class NoReferencesPrivilege(BaseRepackError):
80+
pass
81+
82+
7983
class PostBackfillBatchCallback(typing.Protocol):
8084
def __call__(
8185
self, batch: _introspect.BackfillBatch, /
@@ -824,3 +828,25 @@ def _check_user_permissions(self) -> None:
824828
for table in referring_tables_without_ownership:
825829
message += f"ALTER TABLE {table} OWNER TO {user};\n"
826830
raise NoReferringTableOwnership(message)
831+
832+
referred_fks_without_references_privilege = [
833+
fk
834+
for fk in self.introspector.get_referred_fks(
835+
table=self.table, schema=self.schema
836+
)
837+
if "REFERENCES" not in fk.privileges
838+
]
839+
if referred_fks_without_references_privilege:
840+
user = self.introspector.get_user()
841+
tables_missing_ownership = [
842+
f"{fk.schema}.{fk.name}"
843+
for fk in referred_fks_without_references_privilege
844+
]
845+
message = (
846+
f"Psycopack requires the user to REFERENCES privilege on tables that "
847+
f"{self.schema}.{self.table} has foreign keys to. "
848+
f"You can grant this privilege to your user via the following commands:\n"
849+
)
850+
for table in tables_missing_ownership:
851+
message += f"GRANT REFERENCES ON TABLE {table} TO {user};\n"
852+
raise NoReferencesPrivilege(message)

‎tests/test_repack.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
InvalidPrimaryKeyTypeForConversion,
1414
InvalidStageForReset,
1515
NoCreateAndUsagePrivilegeOnSchema,
16+
NoReferencesPrivilege,
1617
NoReferringTableOwnership,
1718
NotTableOwner,
1819
PrimaryKeyNotFound,
@@ -1785,3 +1786,39 @@ def test_user_without_referring_table_ownership(
17851786
cur=cur,
17861787
schema=schema,
17871788
)
1789+
1790+
1791+
def test_user_without_referred_table_references_privilege(
1792+
connection: _psycopg.Connection,
1793+
) -> None:
1794+
schema = "sweet_schema"
1795+
with connection.cursor() as cur:
1796+
cur.execute(f"CREATE SCHEMA {schema};")
1797+
cur.execute(f"REVOKE CREATE ON SCHEMA {schema} FROM PUBLIC;")
1798+
factories.create_table_for_repacking(
1799+
connection=connection,
1800+
cur=cur,
1801+
table_name="to_repack",
1802+
rows=10,
1803+
schema=schema,
1804+
)
1805+
cur.execute("DROP USER IF EXISTS sweet_user;")
1806+
cur.execute("CREATE USER sweet_user;")
1807+
cur.execute("GRANT CREATE, USAGE ON SCHEMA sweet_schema TO sweet_user;")
1808+
cur.execute("ALTER TABLE sweet_schema.to_repack OWNER TO sweet_user;")
1809+
cur.execute("ALTER TABLE sweet_schema.referring_table OWNER TO sweet_user;")
1810+
cur.execute(
1811+
"ALTER TABLE sweet_schema.not_valid_referring_table OWNER TO sweet_user;"
1812+
)
1813+
cur.execute("SET ROLE sweet_user;")
1814+
with pytest.raises(
1815+
NoReferencesPrivilege,
1816+
match="GRANT REFERENCES ON TABLE sweet_schema.referred_table TO sweet_user;",
1817+
):
1818+
Repack(
1819+
table="to_repack",
1820+
batch_size=1,
1821+
conn=connection,
1822+
cur=cur,
1823+
schema=schema,
1824+
)

0 commit comments

Comments
 (0)