|
6 | 6 | from django.conf import settings |
7 | 7 | from django.db import DEFAULT_DB_ALIAS, OperationalError |
8 | 8 | from freezegun import freeze_time |
| 9 | +from psycopg2 import sql as psycopg2_sql |
9 | 10 | from rest_framework_json_api.serializers import ValidationError |
10 | 11 |
|
11 | 12 | from api.db_utils import ( |
12 | 13 | POSTGRES_TENANT_VAR, |
| 14 | + PostgresEnumMigration, |
13 | 15 | _should_create_index_on_partition, |
14 | 16 | batch_delete, |
15 | 17 | create_objects_in_batches, |
@@ -910,3 +912,61 @@ def test_rls_transaction_cursor_yielded_correctly(self, tenants_fixture): |
910 | 912 | cursor.execute("SELECT 1") |
911 | 913 | result = cursor.fetchone() |
912 | 914 | assert result[0] == 1 |
| 915 | + |
| 916 | + |
| 917 | +class TestPostgresEnumMigration: |
| 918 | + """ |
| 919 | + Verify that PostgresEnumMigration builds DDL statements via psycopg2.sql |
| 920 | + so that enum type names and values are always properly quoted — preventing |
| 921 | + SQL injection through f-string interpolation. |
| 922 | + """ |
| 923 | + |
| 924 | + def _make_mock_schema_editor(self): |
| 925 | + mock_cursor = MagicMock() |
| 926 | + mock_conn = MagicMock() |
| 927 | + mock_conn.cursor.return_value.__enter__ = MagicMock(return_value=mock_cursor) |
| 928 | + mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False) |
| 929 | + mock_schema_editor = MagicMock() |
| 930 | + mock_schema_editor.connection = mock_conn |
| 931 | + return mock_schema_editor, mock_cursor |
| 932 | + |
| 933 | + def test_create_enum_type_generates_correct_sql(self): |
| 934 | + """create_enum_type builds a proper CREATE TYPE … AS ENUM via psycopg2.sql.""" |
| 935 | + migration = PostgresEnumMigration("my_enum", ("val_a", "val_b")) |
| 936 | + schema_editor, mock_cursor = self._make_mock_schema_editor() |
| 937 | + |
| 938 | + migration.create_enum_type(apps=None, schema_editor=schema_editor) |
| 939 | + |
| 940 | + mock_cursor.execute.assert_called_once() |
| 941 | + query_arg = mock_cursor.execute.call_args[0][0] |
| 942 | + assert isinstance( |
| 943 | + query_arg, psycopg2_sql.Composable |
| 944 | + ), "create_enum_type must pass a psycopg2.sql.Composable, not a raw string." |
| 945 | + # Verify the composed SQL structure: CREATE TYPE <Identifier> AS ENUM (<Literals>) |
| 946 | + parts = query_arg.seq |
| 947 | + assert parts[0] == psycopg2_sql.SQL("CREATE TYPE ") |
| 948 | + assert isinstance(parts[1], psycopg2_sql.Identifier) |
| 949 | + assert parts[1].strings == ("my_enum",) |
| 950 | + assert parts[2] == psycopg2_sql.SQL(" AS ENUM (") |
| 951 | + # The enum values are a Composed of Literal items joined by ", " |
| 952 | + enum_literals = [p for p in parts[3].seq if isinstance(p, psycopg2_sql.Literal)] |
| 953 | + assert [lit._wrapped for lit in enum_literals] == ["val_a", "val_b"] |
| 954 | + assert parts[4] == psycopg2_sql.SQL(")") |
| 955 | + |
| 956 | + def test_drop_enum_type_generates_correct_sql(self): |
| 957 | + """drop_enum_type builds a proper DROP TYPE via psycopg2.sql.""" |
| 958 | + migration = PostgresEnumMigration("my_enum", ("val_a",)) |
| 959 | + schema_editor, mock_cursor = self._make_mock_schema_editor() |
| 960 | + |
| 961 | + migration.drop_enum_type(apps=None, schema_editor=schema_editor) |
| 962 | + |
| 963 | + mock_cursor.execute.assert_called_once() |
| 964 | + query_arg = mock_cursor.execute.call_args[0][0] |
| 965 | + assert isinstance( |
| 966 | + query_arg, psycopg2_sql.Composable |
| 967 | + ), "drop_enum_type must pass a psycopg2.sql.Composable, not a raw string." |
| 968 | + # Verify the composed SQL structure: DROP TYPE <Identifier> |
| 969 | + parts = query_arg.seq |
| 970 | + assert parts[0] == psycopg2_sql.SQL("DROP TYPE ") |
| 971 | + assert isinstance(parts[1], psycopg2_sql.Identifier) |
| 972 | + assert parts[1].strings == ("my_enum",) |
0 commit comments