diff --git a/psqlextra/backend/introspection.py b/psqlextra/backend/introspection.py index a9106bd..a82f54c 100644 --- a/psqlextra/backend/introspection.py +++ b/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. diff --git a/psqlextra/backend/schema.py b/psqlextra/backend/schema.py index da7eca6..1e4f2dd 100644 --- a/psqlextra/backend/schema.py +++ b/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: diff --git a/psqlextra/contrib/model_data_migrator.py b/psqlextra/contrib/model_data_migrator.py index 4d90142..5dd640d 100644 --- a/psqlextra/contrib/model_data_migrator.py +++ b/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__ diff --git a/psqlextra/settings.py b/psqlextra/settings.py index b606176..0ea65b2 100644 --- a/psqlextra/settings.py +++ b/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") [ diff --git a/tests/test_model_data_migrator.py b/tests/test_model_data_migrator.py index 2fb2ff2..9cf7569 100644 --- a/tests/test_model_data_migrator.py +++ b/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,) diff --git a/tests/test_settings.py b/tests/test_settings.py index 4451971..f760dda 100644 --- a/tests/test_settings.py +++ b/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"]):