Fix PostgresModelDataMigrator leaving the swapped-in table without privileges, stats and autovacuum

psqlextra/backend/introspection.py

@@ -301,6 +301,33 @@ def get_storage_settings(self, cursor, table_name: str) -> Dict[str, str]:
 
         return storage_settings
 
+    def get_privileges(
+        self, cursor, table_name: str
+    ) -> List[Tuple[Optional[str], str, bool]]:
+        """Gets the privileges granted on the specified table as a list of
+        (grantee, privilege, grantable) tuples.
+
+        The grantee is `None` for privileges granted to PUBLIC.
+        """
+
+        sql = """
+            SELECT
+                CASE WHEN acl.grantee = 0 THEN NULL ELSE pg_catalog.pg_get_userbyid(acl.grantee) END,
+                acl.privilege_type,
+                acl.is_grantable
+            FROM
+                pg_catalog.pg_class c
+            CROSS JOIN LATERAL
+                pg_catalog.aclexplode(c.relacl) acl
+            WHERE
+                c.relname::text = %s
+                AND pg_catalog.pg_table_is_visible(c.oid)
+            ORDER BY 1, 2
+        """
+
+        cursor.execute(sql, (table_name,))
+        return cursor.fetchall()
+
     def get_relations(self, cursor, table_name: str):
         """Gets a dictionary {field_name: (field_name_other_table,
         other_table)} representing all relations in the specified table.

psqlextra/backend/schema.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Any, List, Optional, Type, cast
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
 from unittest import mock
 
 import django
@@ -53,6 +54,8 @@ class PostgresSchemaEditor(SchemaEditor):
     sql_alter_table_storage_setting = "ALTER TABLE %s SET (%s = %s)"
     sql_reset_table_storage_setting = "ALTER TABLE %s RESET (%s)"
 
+    sql_grant_table_privileges = "GRANT %s ON TABLE %s TO %s%s"
+
     sql_alter_table_schema = "ALTER TABLE %s SET SCHEMA %s"
     sql_create_schema = "CREATE SCHEMA %s"
     sql_delete_schema = "DROP SCHEMA %s"
@@ -223,6 +226,65 @@ def clone_model_structure_to_schema(
                 quoted_table_fqn, setting_name, setting_value
             )
 
+    def clone_model_privileges_to_schema(
+        self, model: Type[Model], *, schema_name: str
+    ) -> None:
+        """Grants the privileges granted on the model table to the table that
+        was cloned from it by `clone_model_structure_to_schema`.
+
+        Postgres copies none of them, and the default privileges for
+        the schema the original table is in do not apply to a table
+        created in another schema.
+
+        Run this right before swapping the cloned table in, so that
+        privileges granted or revoked in the meantime are not lost.
+
+        Only privileges on the table itself are copied. These are not:
+
+            - Column privileges, such as `GRANT SELECT (name)`.
+
+            - Privileges on the sequences that
+              `clone_model_structure_to_schema` creates for the
+              cloned table on Django 4.0 and older. Newer versions
+              use identity columns, and inserting into those does
+              not check the privileges on their sequences.
+
+        Arguments:
+            model:
+                Model for which the cloned table was created.
+
+            schema_name:
+                Name of the schema in which the cloned table
+                resides.
+        """
+
+        quoted_table_fqn = f"{self.quote_name(schema_name)}.{self.quote_name(model._meta.db_table)}"
+
+        with self.connection.cursor() as cursor:
+            privileges = self.introspection.get_privileges(
+                cursor, model._meta.db_table
+            )
+
+        privileges_by_grantee: Dict[
+            Tuple[Optional[str], bool], List[str]
+        ] = defaultdict(list)
+        for grantee, privilege, grantable in privileges:
+            privileges_by_grantee[(grantee, grantable)].append(privilege)
+
+        for (
+            grantee,
+            grantable,
+        ), grantee_privileges in privileges_by_grantee.items():
+            self.execute(
+                self.sql_grant_table_privileges
+                % (
+                    ", ".join(grantee_privileges),
+                    quoted_table_fqn,
+                    self.quote_name(grantee) if grantee else "PUBLIC",
+                    " WITH GRANT OPTION" if grantable else "",
+                )
+            )
+
     def clone_model_constraints_and_indexes_to_schema(
         self, model: Type[Model], *, schema_name: str
     ) -> None:

psqlextra/contrib/model_data_migrator.py

@@ -87,14 +87,22 @@ class PostgresModelDataMigrator:
             began. Since the original table is locked, you can
             be sure no more rows are being added or modified.
 
-        9. Move the original table into a backup schema.
+        9. Grant the privileges of the real table on the copy.
+
+            Postgres does not copy them, and the default privileges
+            of the schema the real table is in do not apply to it.
+
+            Column privileges, and on Django 4.0 and older the
+            privileges on the table's sequences, are not copied.
+
+        10. Move the original table into a backup schema.
 
             This allows it to be quickly restored manually
             if the migration is broken in any way.
 
-        10. Move the copied table in place of the real one.
+        11. Move the copied table in place of the real one.
 
-        11. Commit the transaction, which releases the lock.
+        12. Commit the transaction, which releases the lock.
 
     The process is very similiar to how pg_repack rewrites
     an entire table without long-running locks on the table.
@@ -292,25 +300,25 @@ def _migrate_phase_3(self, state: PostgresModelDataMigratorState) -> None:
         # ANALYZE: The table went from 0 to being filled, by running ANALYZE,
         #          we update the statistics, allowing the query planner to
         #          make good decisions.
-        with postgres_prepend_local_search_path(
-            [state.work_schema.name], using=self.using
-        ):
-            self.schema_editor.vacuum_model(self.model, analyze=True)
+        #
+        # VACUUM cannot run in a transaction, which a local search path
+        # needs, so the cloned table is referred to by its schema instead.
+        cloned_table_fqn = self._cloned_table_fqn(state)
+
+        self.schema_editor.vacuum_table(cloned_table_fqn, analyze=True)
 
         # Re-enable autovacuum on the cloned table
-        with postgres_prepend_local_search_path(
-            [state.work_schema.name], using=self.using
-        ):
+        with self.atomic():
             autovacuum_enabled = state.storage_settings.get(
                 "autovacuum_enabled"
             )
             if autovacuum_enabled:
-                self.schema_editor.alter_model_storage_setting(
-                    self.model, "autovacuum_enabled", autovacuum_enabled
+                self.schema_editor.alter_table_storage_setting(
+                    cloned_table_fqn, "autovacuum_enabled", autovacuum_enabled
                 )
             else:
-                self.schema_editor.reset_model_storage_setting(
-                    self.model, "autovacuum_enabled"
+                self.schema_editor.reset_table_storage_setting(
+                    cloned_table_fqn, "autovacuum_enabled"
                 )
 
     def _migrate_phase_4(self, state: PostgresModelDataMigratorState) -> None:
@@ -329,6 +337,12 @@ def _migrate_phase_4(self, state: PostgresModelDataMigratorState) -> None:
                 state.work_schema, state.default_schema
             )
 
+            # Copied last so that privileges granted or revoked while
+            # the cloned table was being filled carry over.
+            self.schema_editor.clone_model_privileges_to_schema(
+                self.model, schema_name=state.work_schema.name
+            )
+
             # Move the original table into the backup schema.
             # Disable autovacuum on it so we don't waste resources
             # keeping it clean.
@@ -347,6 +361,11 @@ def _migrate_phase_4(self, state: PostgresModelDataMigratorState) -> None:
                     self.model, state.default_schema.name
                 )
 
+    def _cloned_table_fqn(self, state: PostgresModelDataMigratorState) -> str:
+        quote_name = self.connection.ops.quote_name
+
+        return f"{quote_name(state.work_schema.name)}.{quote_name(self.table_name)}"
+
     @property
     def model_name(self) -> str:
         return self.model.__name__

psqlextra/settings.py

@@ -94,6 +94,11 @@ def postgres_prepend_local_search_path(
 
     connection = connections[using]
 
+    if not connection.in_atomic_block:
+        raise SuspiciousOperation(
+            "SET LOCAL makes no sense outside a transaction. Start a transaction first."
+        )
+
     with connection.cursor() as cursor:
         cursor.execute("SHOW search_path")
         [

tests/test_model_data_migrator.py

@@ -1,3 +1,5 @@
+import os
+
 from datetime import timedelta
 from unittest.mock import MagicMock
 
@@ -10,6 +12,7 @@
 from psqlextra.schema import PostgresSchema
 from psqlextra.settings import postgres_prepend_local_search_path
 
+from . import db_introspection
 from .fake_model import delete_fake_model, get_fake_model
 
 pytestmark = pytest.mark.skipif(
@@ -27,7 +30,24 @@ def fake_model():
     delete_fake_model(model)
 
 
-def _create_migrator(model, *, keep_backup_schema=True, fail=False):
+@pytest.fixture
+def role():
+    name = f"psqlextra_{os.urandom(4).hex()}"
+    quoted_name = connection.ops.quote_name(name)
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"CREATE ROLE {quoted_name}")
+
+    yield name
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"DROP OWNED BY {quoted_name}")
+        cursor.execute(f"DROP ROLE {quoted_name}")
+
+
+def _create_migrator(
+    model, *, keep_backup_schema=True, fail=False, while_filling=None
+):
     class Migrator(PostgresModelDataMigrator):
         operation_timeout = timedelta(minutes=1)
 
@@ -36,6 +56,9 @@ def fill_cloned_table_lockless(self, work_schema, default_schema):
                 with postgres_prepend_local_search_path([work_schema.name]):
                     self.model.objects.create(name="new")
 
+            if while_filling:
+                while_filling()
+
             if fail:
                 raise RuntimeError("fill failed")
 
@@ -108,3 +131,101 @@ def test_model_data_migrator_deletes_its_schemas_when_it_fails(fake_model):
 
     assert _list_schemas_of(fake_model) == []
     assert list(fake_model.objects.values_list("name", flat=True)) == ["old"]
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_keeps_the_storage_settings(fake_model):
+    with connection.schema_editor() as schema_editor:
+        schema_editor.alter_model_storage_setting(
+            fake_model, "fillfactor", "80"
+        )
+
+    _create_migrator(fake_model, keep_backup_schema=False).migrate()
+
+    with transaction.atomic():
+        assert db_introspection.get_storage_settings(
+            fake_model._meta.db_table
+        ) == {"fillfactor": "80"}
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_analyzes_the_swapped_in_table(fake_model):
+    _create_migrator(fake_model, keep_backup_schema=False).migrate()
+
+    with connection.cursor() as cursor:
+        cursor.execute(
+            "SELECT attname FROM pg_stats WHERE schemaname = current_schema() AND tablename = %s",
+            (fake_model._meta.db_table,),
+        )
+        assert "name" in {attname for attname, in cursor.fetchall()}
+
+
+def _list_privileges_of(model):
+    with connection.cursor() as cursor:
+        cursor.execute(
+            "SELECT relacl::text[] FROM pg_class WHERE oid = %s::regclass",
+            (connection.ops.quote_name(model._meta.db_table),),
+        )
+        [acl] = cursor.fetchone()
+
+    return set(acl or [])
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_keeps_the_privileges(fake_model, role):
+    quoted_table_name = connection.ops.quote_name(fake_model._meta.db_table)
+    quoted_role_name = connection.ops.quote_name(role)
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"GRANT SELECT ON {quoted_table_name} TO PUBLIC")
+        cursor.execute(
+            f"GRANT SELECT, DELETE ON {quoted_table_name} TO {quoted_role_name}"
+        )
+        cursor.execute(
+            f"GRANT UPDATE ON {quoted_table_name} TO {quoted_role_name} WITH GRANT OPTION"
+        )
+
+    privileges = _list_privileges_of(fake_model)
+
+    _create_migrator(fake_model, keep_backup_schema=False).migrate()
+
+    assert _list_privileges_of(fake_model) == privileges
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_keeps_the_autovacuum_setting(fake_model):
+    with connection.schema_editor() as schema_editor:
+        schema_editor.alter_model_storage_setting(
+            fake_model, "autovacuum_enabled", "true"
+        )
+
+    _create_migrator(fake_model, keep_backup_schema=False).migrate()
+
+    with transaction.atomic():
+        assert db_introspection.get_storage_settings(
+            fake_model._meta.db_table
+        ) == {"autovacuum_enabled": "true"}
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_keeps_the_privileges_granted_while_filling(
+    fake_model, role
+):
+    quoted_table_name = connection.ops.quote_name(fake_model._meta.db_table)
+
+    def grant():
+        with connection.cursor() as cursor:
+            cursor.execute(
+                f"GRANT SELECT ON {quoted_table_name} TO {connection.ops.quote_name(role)}"
+            )
+
+    _create_migrator(
+        fake_model, keep_backup_schema=False, while_filling=grant
+    ).migrate()
+
+    with connection.cursor() as cursor:
+        cursor.execute(
+            "SELECT has_table_privilege(%s, %s, 'SELECT')",
+            (role, quoted_table_name),
+        )
+        assert cursor.fetchone() == (True,)

tests/test_settings.py

@@ -80,6 +80,13 @@ def test_postgres_prepend_local_search_path():
     assert _get_current_setting("search_path") == '"$user", public'
 
 
+@pytest.mark.django_db(transaction=True)
+def test_postgres_prepend_local_search_path_no_transaction():
+    with pytest.raises(SuspiciousOperation):
+        with postgres_prepend_local_search_path(["a"]):
+            pass
+
+
 def test_postgres_prepend_local_search_path_nested():
     with postgres_prepend_local_search_path(["a", "b"]):
         with postgres_prepend_local_search_path(["c"]):