Merge pull request #30 from kraken-tech/forbid-user-without-access

marcelofern · web-flow · commit 41ebf9b2e317 · 2025-05-15T09:02:41.000+12:00
Stop user without access permissions
diff --git a/src/psycopack/__init__.py b/src/psycopack/__init__.py
@@ -10,6 +10,10 @@
     InvalidIndexes,
     InvalidPrimaryKeyTypeForConversion,
     InvalidStageForReset,
+    NoCreateAndUsagePrivilegeOnSchema,
+    NoReferencesPrivilege,
+    NoReferringTableOwnership,
+    NotTableOwner,
     PrimaryKeyNotFound,
     ReferringForeignKeyInDifferentSchema,
     Repack,
@@ -30,6 +34,10 @@
     "InvalidIndexes",
     "InvalidPrimaryKeyTypeForConversion",
     "InvalidStageForReset",
+    "NoCreateAndUsagePrivilegeOnSchema",
+    "NoReferencesPrivilege",
+    "NoReferringTableOwnership",
+    "NotTableOwner",
     "PrimaryKeyNotFound",
     "ReferringForeignKeyInDifferentSchema",
     "Repack",
diff --git a/src/psycopack/_introspect.py b/src/psycopack/_introspect.py
@@ -30,6 +30,14 @@ class ReferringForeignKey:
     is_validated: bool
     referring_table: str
     schema: str
+    is_owned_by_user: bool
+
+
+@dataclasses.dataclass
+class ReferredForeignKey:
+    name: str
+    schema: str
+    privileges: list[str]
 
 
 @dataclasses.dataclass
@@ -239,7 +247,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 +277,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 +292,74 @@ 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 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:
@@ -529,3 +615,48 @@ def get_backfill_batch(self, *, table: str) -> BackfillBatch | None:
         if not (result := self.cur.fetchone()):
             return None
         return BackfillBatch(id=result[0], start=result[1], end=result[2])
+
+    def get_user(self) -> str:
+        self.cur.execute(psycopg.sql.SQL("SELECT current_user;").as_string(self.conn))
+        result = self.cur.fetchone()
+        assert result is not None
+        return str(result[0])
+
+    def has_create_and_usage_privilege_on_schema(self) -> bool:
+        self.cur.execute(
+            psycopg.sql.SQL(
+                dedent("""
+                SELECT
+                  has_schema_privilege({schema}, 'CREATE'),
+                  has_schema_privilege({schema}, 'USAGE');
+                """)
+            )
+            .format(schema=psycopg.sql.Literal(self.schema))
+            .as_string(self.conn)
+        )
+        result = self.cur.fetchone()
+        assert result is not None
+        return bool(result[0] and result[1])
+
+    def is_table_owner(self, *, table: str, schema: str) -> bool:
+        self.cur.execute(
+            psycopg.sql.SQL(
+                dedent("""
+                SELECT
+                  tableowner = current_user
+                FROM
+                  pg_tables
+                WHERE
+                  schemaname = {schema}
+                  AND tablename = {table}
+                """)
+            )
+            .format(
+                schema=psycopg.sql.Literal(schema),
+                table=psycopg.sql.Literal(table),
+            )
+            .as_string(self.conn)
+        )
+        result = self.cur.fetchone()
+        assert result is not None
+        return bool(result[0])
diff --git a/src/psycopack/_repack.py b/src/psycopack/_repack.py
@@ -64,6 +64,22 @@ class ReferringForeignKeyInDifferentSchema(BaseRepackError):
     pass
 
 
+class NoCreateAndUsagePrivilegeOnSchema(BaseRepackError):
+    pass
+
+
+class NotTableOwner(BaseRepackError):
+    pass
+
+
+class NoReferringTableOwnership(BaseRepackError):
+    pass
+
+
+class NoReferencesPrivilege(BaseRepackError):
+    pass
+
+
 class PostBackfillBatchCallback(typing.Protocol):
     def __call__(
         self, batch: _introspect.BackfillBatch, /
@@ -175,6 +191,7 @@ def __init__(
         )
         self._pk_column = ""
         self.schema = schema
+        self._check_user_permissions()
 
     @property
     def pk_column(self) -> str:
@@ -776,3 +793,60 @@ def _create_referring_fks(self) -> None:
                 self.command.validate_constraint(
                     table=fk.referring_table, constraint=constraint_name
                 )
+
+    def _check_user_permissions(self) -> None:
+        if not self.introspector.has_create_and_usage_privilege_on_schema():
+            user = self.introspector.get_user()
+            raise NoCreateAndUsagePrivilegeOnSchema(
+                f"Psycopack requires the database user to have CREATE and USAGE "
+                f"privilege on the {self.schema} schema to create auxiliary "
+                f"objects. You can grant it to your user via:\n"
+                f"GRANT CREATE, USAGE ON SCHEMA {self.schema} TO {user};"
+            )
+
+        if not self.introspector.is_table_owner(table=self.table, schema=self.schema):
+            user = self.introspector.get_user()
+            raise NotTableOwner(
+                f"Psycopack requires the database user to have ownership of "
+                f"the table {self.schema}.{self.table}. You can grant it to "
+                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)
+
+        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)
diff --git a/tests/test_repack.py b/tests/test_repack.py
