Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add support to schema editor for moving tables between schemas
- Loading branch information
Showing
8 changed files
with
361 additions
and
80 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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)] |
Oops, something went wrong.