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.

Author: marcelofern

src/psycopack/__init__.py

@@ -11,6 +11,7 @@
     InvalidPrimaryKeyTypeForConversion,
     InvalidStageForReset,
     NoCreateAndUsagePrivilegeOnSchema,
+    NoReferencesPrivilege,
     NoReferringTableOwnership,
     NotTableOwner,
     PrimaryKeyNotFound,
@@ -34,6 +35,7 @@
     "InvalidPrimaryKeyTypeForConversion",
     "InvalidStageForReset",
     "NoCreateAndUsagePrivilegeOnSchema",
+    "NoReferencesPrivilege",
     "NoReferringTableOwnership",
     "NotTableOwner",
     "PrimaryKeyNotFound",

src/psycopack/_introspect.py

@@ -33,6 +33,13 @@ class ReferringForeignKey:
     is_owned_by_user: bool
 
 
+@dataclasses.dataclass
+class ReferredForeignKey:
+    name: str
+    schema: str
+    privileges: list[str]
+
+
 @dataclasses.dataclass
 class Trigger:
     name: str
@@ -290,6 +297,71 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
             for name, definition, is_validated, referring_table, schema, is_owned_by_user in results
         ]
 
+    def get_referred_fks(self, *, table: str, schema: str) -> list[ReferredForeignKey]:
+        self.cur.execute(
+            psycopg.sql.SQL(
+                dedent("""
+                SELECT
+                  referred_pg_class.relname AS referred_table,
+                  referred_pg_namespace.nspname AS referred_schema,
+                  array_agg(
+                    table_privileges.privilege_type
+                    ORDER BY table_privileges.privilege_type
+                  ) AS privileges
+                FROM
+                  pg_catalog.pg_constraint
+                INNER JOIN
+                  pg_catalog.pg_class AS referring_pg_class
+                  ON (pg_constraint.conrelid = referring_pg_class.oid)
+                INNER JOIN
+                  pg_catalog.pg_class AS referred_pg_class
+                  ON (pg_constraint.confrelid = referred_pg_class.oid)
+                INNER JOIN
+                  pg_catalog.pg_namespace AS referring_pg_namespace
+                  ON (referring_pg_namespace.oid = referring_pg_class.relnamespace)
+                INNER JOIN
+                  pg_catalog.pg_namespace AS referred_pg_namespace
+                  ON (referred_pg_namespace.oid = referred_pg_class.relnamespace)
+                LEFT OUTER JOIN
+                  information_schema.table_privileges
+                  ON (
+                    table_privileges.table_name = referred_pg_class.relname
+                    AND table_privileges.table_schema = referred_pg_namespace.nspname
+                  )
+                WHERE
+                  pg_constraint.conrelid = referring_pg_class.oid
+                  AND referring_pg_class.relname = {table}
+                  AND referring_pg_namespace.nspname = {schema}
+                  AND pg_constraint.contype = 'f'
+                GROUP BY
+                  referred_pg_namespace.nspname,
+                  referred_pg_class.relname
+                ORDER BY
+                  referred_pg_namespace.nspname,
+                  referred_pg_class.relname;
+                """)
+            )
+            .format(
+                table=psycopg.sql.Literal(table),
+                schema=psycopg.sql.Literal(schema),
+            )
+            .as_string(self.conn)
+        )
+        results = self.cur.fetchall()
+        assert results is not None
+        return [
+            ReferredForeignKey(
+                name=name,
+                schema=schema,
+                privileges=(
+                    privileges.strip("{}").split(",")
+                    if "NULL" not in privileges
+                    else []
+                ),
+            )
+            for name, schema, privileges in results
+        ]
+
     def table_is_empty(self, *, table: str) -> int:
         self.cur.execute(
             psycopg.sql.SQL(

src/psycopack/_repack.py

@@ -76,6 +76,10 @@ class NoReferringTableOwnership(BaseRepackError):
     pass
 
 
+class NoReferencesPrivilege(BaseRepackError):
+    pass
+
+
 class PostBackfillBatchCallback(typing.Protocol):
     def __call__(
         self, batch: _introspect.BackfillBatch, /
@@ -824,3 +828,25 @@ def _check_user_permissions(self) -> None:
             for table in referring_tables_without_ownership:
                 message += f"ALTER TABLE {table} OWNER TO {user};\n"
             raise NoReferringTableOwnership(message)
+
+        referred_fks_without_references_privilege = [
+            fk
+            for fk in self.introspector.get_referred_fks(
+                table=self.table, schema=self.schema
+            )
+            if "REFERENCES" not in fk.privileges
+        ]
+        if referred_fks_without_references_privilege:
+            user = self.introspector.get_user()
+            tables_missing_ownership = [
+                f"{fk.schema}.{fk.name}"
+                for fk in referred_fks_without_references_privilege
+            ]
+            message = (
+                f"Psycopack requires the user to REFERENCES privilege on tables that "
+                f"{self.schema}.{self.table} has foreign keys to. "
+                f"You can grant this privilege to your user via the following commands:\n"
+            )
+            for table in tables_missing_ownership:
+                message += f"GRANT REFERENCES ON TABLE {table} TO {user};\n"
+            raise NoReferencesPrivilege(message)

tests/test_repack.py

@@ -13,6 +13,7 @@
     InvalidPrimaryKeyTypeForConversion,
     InvalidStageForReset,
     NoCreateAndUsagePrivilegeOnSchema,
+    NoReferencesPrivilege,
     NoReferringTableOwnership,
     NotTableOwner,
     PrimaryKeyNotFound,
@@ -1785,3 +1786,39 @@ def test_user_without_referring_table_ownership(
                 cur=cur,
                 schema=schema,
             )
+
+
+def test_user_without_referred_table_references_privilege(
+    connection: _psycopg.Connection,
+) -> None:
+    schema = "sweet_schema"
+    with connection.cursor() as cur:
+        cur.execute(f"CREATE SCHEMA {schema};")
+        cur.execute(f"REVOKE CREATE ON SCHEMA {schema} FROM PUBLIC;")
+        factories.create_table_for_repacking(
+            connection=connection,
+            cur=cur,
+            table_name="to_repack",
+            rows=10,
+            schema=schema,
+        )
+        cur.execute("DROP USER IF EXISTS sweet_user;")
+        cur.execute("CREATE USER sweet_user;")
+        cur.execute("GRANT CREATE, USAGE ON SCHEMA sweet_schema TO sweet_user;")
+        cur.execute("ALTER TABLE sweet_schema.to_repack OWNER TO sweet_user;")
+        cur.execute("ALTER TABLE sweet_schema.referring_table OWNER TO sweet_user;")
+        cur.execute(
+            "ALTER TABLE sweet_schema.not_valid_referring_table OWNER TO sweet_user;"
+        )
+        cur.execute("SET ROLE sweet_user;")
+        with pytest.raises(
+            NoReferencesPrivilege,
+            match="GRANT REFERENCES ON TABLE sweet_schema.referred_table TO sweet_user;",
+        ):
+            Repack(
+                table="to_repack",
+                batch_size=1,
+                conn=connection,
+                cur=cur,
+                schema=schema,
+            )