Add support to schema editor for moving tables between schemas

Photonios · Photonios · commit 6eff3f12a1c2 · 2023-04-11T09:48:08.000+03:00
diff --git a/psqlextra/backend/introspection.py b/psqlextra/backend/introspection.py
@@ -1,9 +1,6 @@
-from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple
 
-from django.db import transaction
-
 from psqlextra.types import PostgresPartitioningMethod
 
 from . import base_impl
@@ -250,6 +247,7 @@ def get_storage_settings(self, cursor, table_name: str) -> Dict[str, str]:
                 pg_catalog.pg_am am ON (c.relam = am.oid)
             WHERE
                 c.relname::text = %s
+                AND pg_catalog.pg_table_is_visible(c.oid)
         """
 
         cursor.execute(sql, (table_name,))
@@ -288,29 +286,3 @@ def get_relations(self, cursor, table_name: str):
             [table_name],
         )
         return {row[0]: (row[2], row[1]) for row in cursor.fetchall()}
-
-    @contextmanager
-    def in_search_path(self, search_path: List[str]):
-        """Changes the Postgres `search_path` within the context and switches
-        it back when it exits."""
-
-        # Wrap in a transaction so a savepoint is created. If
-        # something goes wrong, the `SET LOCAL search_path`
-        # statement will be rolled back.
-        with transaction.atomic(using=self.connection.alias):
-            with self.connection.cursor() as cursor:
-                cursor.execute("SHOW search_path")
-                (original_search_path,) = cursor.fetchone()
-
-                # Syntax in Postgres is a bit weird here. It isn't really
-                # a list of names like in `WHERE bla in (val1, val2)`.
-                placeholder = ", ".join(["%s" for _ in search_path])
-                cursor.execute(
-                    f"SET LOCAL search_path = {placeholder}", search_path
-                )
-
-                yield self
-
-                cursor.execute(
-                    f"SET LOCAL search_path = {original_search_path}"
-                )
diff --git a/psqlextra/backend/operations.py b/psqlextra/backend/operations.py
@@ -21,6 +21,3 @@ class PostgresOperations(base_impl.operations()):
         SQLUpdateCompiler,
         SQLInsertCompiler,
     ]
-
-    def default_schema_name(self) -> str:
-        return "public"
diff --git a/psqlextra/backend/schema.py b/psqlextra/backend/schema.py
@@ -12,6 +12,10 @@
 from django.db.backends.ddl_references import Statement
 from django.db.models import Field, Model
 
+from psqlextra.settings import (
+    postgres_prepend_local_search_path,
+    postgres_reset_local_search_path,
+)
 from psqlextra.type_assertions import is_sql_with_params
 from psqlextra.types import PostgresPartitioningMethod
 
@@ -40,6 +44,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_alter_table_schema = "ALTER TABLE %s SET SCHEMA %s"
+
     sql_create_view = "CREATE VIEW %s AS (%s)"
     sql_replace_view = "CREATE OR REPLACE VIEW %s AS (%s)"
     sql_drop_view = "DROP VIEW IF EXISTS %s"
@@ -203,8 +209,8 @@ def clone_model_constraints_and_indexes_to_schema(
                 resides.
         """
 
-        with self.introspection.in_search_path(
-            [schema_name, self.connection.ops.default_schema_name()]
+        with postgres_prepend_local_search_path(
+            [schema_name], using=self.connection.alias
         ):
             for constraint in model._meta.constraints:
                 self.add_constraint(model, constraint)
@@ -226,8 +232,8 @@ def clone_model_constraints_and_indexes_to_schema(
                 # Django creates primary keys later added to the model with
                 # a custom name. We want the name as it was created originally.
                 if field.primary_key:
-                    with self.introspection.in_search_path(
-                        [self.connection.ops.default_schema_name()]
+                    with postgres_reset_local_search_path(
+                        using=self.connection.alias
                     ):
                         [primary_key_name] = self._constraint_names(
                             model, primary_key=True
@@ -251,8 +257,8 @@ def clone_model_constraints_and_indexes_to_schema(
                 # a separate transaction later to validate the entries without
                 # acquiring a AccessExclusiveLock.
                 if field.remote_field:
-                    with self.introspection.in_search_path(
-                        [self.connection.ops.default_schema_name()]
+                    with postgres_reset_local_search_path(
+                        using=self.connection.alias
                     ):
                         [fk_name] = self._constraint_names(
                             model, [field.column], foreign_key=True
@@ -277,8 +283,8 @@ def clone_model_constraints_and_indexes_to_schema(
                 # manually.
                 field_check = field.db_parameters(self.connection).get("check")
                 if field_check:
-                    with self.introspection.in_search_path(
-                        [self.connection.ops.default_schema_name()]
+                    with postgres_reset_local_search_path(
+                        using=self.connection.alias
                     ):
                         [field_check_name] = self._constraint_names(
                             model,
@@ -337,10 +343,12 @@ def clone_model_foreign_keys_to_schema(
                 resides.
         """
 
-        with self.introspection.in_search_path(
-            [schema_name, self.connection.ops.default_schema_name()]
+        constraint_names = self._constraint_names(model, foreign_key=True)
+
+        with postgres_prepend_local_search_path(
+            [schema_name], using=self.connection.alias
         ):
-            for fk_name in self._constraint_names(model, foreign_key=True):
+            for fk_name in constraint_names:
                 self.execute(
                     self.sql_validate_fk
                     % (
@@ -438,6 +446,47 @@ def reset_model_storage_setting(
 
         self.reset_table_storage_setting(model._meta.db_table, name)
 
+    def alter_table_schema(self, table_name: str, schema_name: str) -> None:
+        """Moves the specified table into the specified schema.
+
+        WARNING: Moving models into a different schema than the default
+        will break querying the model.
+
+        Arguments:
+            table_name:
+                Name of the table to move into the specified schema.
+
+            schema_name:
+                Name of the schema to move the table to.
+        """
+
+        self.execute(
+            self.sql_alter_table_schema
+            % (self.quote_name(table_name), self.quote_name(schema_name))
+        )
+
+    def alter_model_schema(self, model: Type[Model], schema_name: str) -> None:
+        """Moves the specified model's table into the specified schema.
+
+        WARNING: Moving models into a different schema than the default
+        will break querying the model.
+
+        Arguments:
+            model:
+               Model of which to move the table.
+
+            schema_name:
+                Name of the schema to move the model's table to.
+        """
+
+        self.execute(
+            self.sql_alter_table_schema
+            % (
+                self.quote_name(model._meta.db_table),
+                self.quote_name(schema_name),
+            )
+        )
+
     def refresh_materialized_view_model(
         self, model: Type[Model], concurrently: bool = False
     ) -> None:
diff --git a/psqlextra/settings.py b/psqlextra/settings.py
@@ -0,0 +1,118 @@
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Union
+
+from django.core.exceptions import SuspiciousOperation
+from django.db import DEFAULT_DB_ALIAS, connections
+
+
+@contextmanager
+def postgres_set_local(
+    *,
+    using: str = DEFAULT_DB_ALIAS,
+    **options: Dict[str, Optional[Union[str, int, float, List[str]]]],
+) -> None:
+    """Sets the specified PostgreSQL options using SET LOCAL so that they apply
+    to the current transacton only.
+
+    The effect is undone when the context manager exits.
+
+    See https://www.postgresql.org/docs/current/runtime-config-client.html
+    for an overview of all available options.
+    """
+
+    connection = connections[using]
+    qn = connection.ops.quote_name
+
+    if not connection.in_atomic_block:
+        raise SuspiciousOperation(
+            "SET LOCAL makes no sense outside a transaction. Start a transaction first."
+        )
+
+    sql = []
+    params = []
+    for name, value in options.items():
+        if value is None:
+            sql.append(f"SET LOCAL {qn(name)} TO DEFAULT")
+            continue
+
+        # Settings that accept a list of values are actually
+        # stored as string lists. We cannot just pass a list
+        # of values. We have to create the comma separated
+        # string ourselves.
+        if isinstance(value, list) or isinstance(value, tuple):
+            placeholder = ", ".join(["%s" for _ in value])
+            params.extend(value)
+        else:
+            placeholder = "%s"
+            params.append(value)
+
+        sql.append(f"SET LOCAL {qn(name)} = {placeholder}")
+
+    with connection.cursor() as cursor:
+        cursor.execute(
+            "SELECT name, setting FROM pg_settings WHERE name = ANY(%s)",
+            (list(options.keys()),),
+        )
+        original_values = dict(cursor.fetchall())
+        cursor.execute("; ".join(sql), params)
+
+    yield
+
+    # Put everything back to how it was. DEFAULT is
+    # not good enough as a outer SET LOCAL might
+    # have set a different value.
+    with connection.cursor() as cursor:
+        sql = []
+        params = []
+
+        for name, value in options.items():
+            original_value = original_values.get(name)
+            if original_value:
+                sql.append(f"SET LOCAL {qn(name)} = {original_value}")
+            else:
+                sql.append(f"SET LOCAL {qn(name)} TO DEFAULT")
+
+        cursor.execute("; ".join(sql), params)
+
+
+@contextmanager
+def postgres_set_local_search_path(
+    search_path: List[str], *, using: str = DEFAULT_DB_ALIAS
+) -> None:
+    """Sets the search path to the specified schemas."""
+
+    with postgres_set_local(search_path=search_path, using=using):
+        yield
+
+
+@contextmanager
+def postgres_prepend_local_search_path(
+    search_path: List[str], *, using: str = DEFAULT_DB_ALIAS
+) -> None:
+    """Prepends the current local search path with the specified schemas."""
+
+    connection = connections[using]
+
+    with connection.cursor() as cursor:
+        cursor.execute("SHOW search_path")
+        [
+            original_search_path,
+        ] = cursor.fetchone()
+
+        placeholders = ", ".join(["%s" for _ in search_path])
+        cursor.execute(
+            f"SET LOCAL search_path = {placeholders}, {original_search_path}",
+            tuple(search_path),
+        )
+
+        yield
+
+        cursor.execute(f"SET LOCAL search_path = {original_search_path}")
+
+
+@contextmanager
+def postgres_reset_local_search_path(*, using: str = DEFAULT_DB_ALIAS) -> None:
+    """Resets the local search path to the default."""
+
+    with postgres_set_local(search_path=None, using=using):
+        yield
diff --git a/tests/db_introspection.py b/tests/db_introspection.py
@@ -9,15 +9,14 @@
 
 from django.db import connection
 
+from psqlextra.settings import postgres_set_local
+
 
 @contextmanager
 def introspect(schema_name: Optional[str] = None):
-    default_schema_name = connection.ops.default_schema_name()
-    search_path = [schema_name or default_schema_name]
-
-    with connection.introspection.in_search_path(search_path) as introspection:
+    with postgres_set_local(search_path=schema_name or None):
         with connection.cursor() as cursor:
-            yield introspection, cursor
+            yield connection.introspection, cursor
 
 
 def table_names(
diff --git a/tests/test_schema_editor_alter_schema.py b/tests/test_schema_editor_alter_schema.py
@@ -0,0 +1,44 @@
+import pytest
+
+from django.db import connection, models
+
+from psqlextra.backend.schema import PostgresSchemaEditor
+
+from .fake_model import get_fake_model
+
+
+@pytest.fixture
+def fake_model():
+    return get_fake_model(
+        {
+            "text": models.TextField(),
+        }
+    )
+
+
+def test_schema_editor_alter_table_schema(fake_model):
+    obj = fake_model.objects.create(text="hello")
+
+    with connection.cursor() as cursor:
+        cursor.execute("CREATE SCHEMA target")
+
+    schema_editor = PostgresSchemaEditor(connection)
+    schema_editor.alter_table_schema(fake_model._meta.db_table, "target")
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"SELECT * FROM target.{fake_model._meta.db_table}")
+        assert cursor.fetchall() == [(obj.id, obj.text)]
+
+
+def test_schema_editor_alter_model_schema(fake_model):
+    obj = fake_model.objects.create(text="hello")
+
+    with connection.cursor() as cursor:
+        cursor.execute("CREATE SCHEMA target")
+
+    schema_editor = PostgresSchemaEditor(connection)
+    schema_editor.alter_model_schema(fake_model, "target")
+
+    with connection.cursor() as cursor:
+        cursor.execute(f"SELECT * FROM target.{fake_model._meta.db_table}")
+        assert cursor.fetchall() == [(obj.id, obj.text)]
diff --git a/tests/test_schema_editor_clone_model_to_schema.py b/tests/test_schema_editor_clone_model_to_schema.py
diff --git a/tests/test_settings.py b/tests/test_settings.py

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,3 @@ class PostgresOperations(base_impl.operations()):`
`21`	`21`	`SQLUpdateCompiler,`
`22`	`22`	`SQLInsertCompiler,`
`23`	`23`	`]`
`24`		`-`
`25`		`- def default_schema_name(self) -> str:`
`26`		`- return "public"`