User must be OWNER of referring tables

Prior to this change, there was no check to verify that the user was an
OWNER of the referring tables (tables that have fks pointing to the
table to be repacked).

This is a problem because Psycopack needs to create new foreign keys on
that table that point to the shadow table, and that is not possible
unless the user is the table owner.

This change adds an explicit check to ensure the user is the owner of
the referring tables before proceeding.

Author: marcelofern

src/psycopack/__init__.py

@@ -11,6 +11,7 @@
     InvalidPrimaryKeyTypeForConversion,
     InvalidStageForReset,
     NoCreateAndUsagePrivilegeOnSchema,
+    NoReferringTableOwnership,
     NotTableOwner,
     PrimaryKeyNotFound,
     ReferringForeignKeyInDifferentSchema,
@@ -33,6 +34,7 @@
     "InvalidPrimaryKeyTypeForConversion",
     "InvalidStageForReset",
     "NoCreateAndUsagePrivilegeOnSchema",
+    "NoReferringTableOwnership",
     "NotTableOwner",
     "PrimaryKeyNotFound",
     "ReferringForeignKeyInDifferentSchema",

src/psycopack/_introspect.py

@@ -30,6 +30,7 @@ class ReferringForeignKey:
     is_validated: bool
     referring_table: str
     schema: str
+    is_owned_by_user: bool
 
 
 @dataclasses.dataclass
@@ -239,7 +240,16 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
                   pg_get_constraintdef(pg_constraint.oid) AS definition,
                   pg_constraint.convalidated AS is_validated,
                   referring_pg_class.relname AS referring_table,
-                  pg_namespace.nspname AS schema
+                  pg_namespace.nspname AS schema,
+                  (
+                    SELECT
+                      tableowner = current_user
+                    FROM
+                      pg_tables
+                    WHERE
+                      schemaname = {schema}
+                      AND tablename = referring_pg_class.relname
+                  ) AS is_owned_by_user
                 FROM
                   pg_catalog.pg_constraint
                 INNER JOIN
@@ -260,7 +270,10 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
                   definition;
                 """)
             )
-            .format(table=psycopg.sql.Literal(table))
+            .format(
+                table=psycopg.sql.Literal(table),
+                schema=psycopg.sql.Literal(self.schema),
+            )
             .as_string(self.conn)
         )
         results = self.cur.fetchall()
@@ -272,8 +285,9 @@ def get_referring_fks(self, *, table: str) -> list[ReferringForeignKey]:
                 is_validated=is_validated,
                 referring_table=referring_table,
                 schema=schema,
+                is_owned_by_user=is_owned_by_user,
             )
-            for name, definition, is_validated, referring_table, schema in results
+            for name, definition, is_validated, referring_table, schema, is_owned_by_user in results
         ]
 
     def table_is_empty(self, *, table: str) -> int:

src/psycopack/_repack.py

@@ -72,6 +72,10 @@ class NotTableOwner(BaseRepackError):
     pass
 
 
+class NoReferringTableOwnership(BaseRepackError):
+    pass
+
+
 class PostBackfillBatchCallback(typing.Protocol):
     def __call__(
         self, batch: _introspect.BackfillBatch, /
@@ -804,3 +808,19 @@ def _check_user_permissions(self) -> None:
                 f"your user via:\n"
                 f"ALTER TABLE {self.schema}.{self.table} OWNER TO {user};"
             )
+
+        referring_tables_without_ownership = [
+            f"{fk.schema}.{fk.referring_table}"
+            for fk in self.introspector.get_referring_fks(table=self.table)
+            if not fk.is_owned_by_user
+        ]
+        if referring_tables_without_ownership:
+            user = self.introspector.get_user()
+            message = (
+                f"Psycopack requires the user to have ownership of tables that "
+                f"have foreign keys pointing to the table {self.schema}.{self.table}. "
+                f"You can grant ownership to your user via the following commands:\n"
+            )
+            for table in referring_tables_without_ownership:
+                message += f"ALTER TABLE {table} OWNER TO {user};\n"
+            raise NoReferringTableOwnership(message)

tests/test_repack.py

@@ -13,6 +13,7 @@
     InvalidPrimaryKeyTypeForConversion,
     InvalidStageForReset,
     NoCreateAndUsagePrivilegeOnSchema,
+    NoReferringTableOwnership,
     NotTableOwner,
     PrimaryKeyNotFound,
     ReferringForeignKeyInDifferentSchema,
@@ -1752,3 +1753,35 @@ def test_user_without_table_ownership(
                 cur=cur,
                 schema=schema,
             )
+
+
+def test_user_without_referring_table_ownership(
+    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("SET ROLE sweet_user;")
+        with pytest.raises(
+            NoReferringTableOwnership,
+            match="ALTER TABLE sweet_schema.referring_table OWNER TO sweet_user;",
+        ):
+            Repack(
+                table="to_repack",
+                batch_size=1,
+                conn=connection,
+                cur=cur,
+                schema=schema,
+            )