diff --git a/CHANGES.rst b/CHANGES.rst index f7ee583..3fc8ebc 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -14,6 +14,7 @@ Features: * #127 Drop Python-3.7 support. * #83 Drop Django-2.2 support. * #134 Support adding COLUMN with UNIQUE; adding column without UNIQUE then add UNIQUE CONSTRAINT. +* #135 Support adding BinaryField. Bug Fixes: diff --git a/django_redshift_backend/base.py b/django_redshift_backend/base.py index 508a8c7..ae572ab 100644 --- a/django_redshift_backend/base.py +++ b/django_redshift_backend/base.py @@ -35,7 +35,9 @@ from django.db.utils import NotSupportedError, ProgrammingError from django_redshift_backend.meta import DistKey, SortKey +from psycopg2.extensions import Binary +from .psycopg2adapter import RedshiftBinary logger = logging.getLogger("django.db.backends") @@ -154,6 +156,10 @@ def _get_type_default(field): return default +def _remove_length_from_type(column_type): + return re.sub(r"\(.*", "", column_type) + + class DatabaseSchemaEditor(BasePGDatabaseSchemaEditor): sql_create_table = "CREATE TABLE %(table)s (%(definition)s) %(options)s" sql_delete_fk = "ALTER TABLE %(table)s DROP CONSTRAINT %(name)s" @@ -187,6 +193,25 @@ def remove_index(self, model, index, concurrently=False): # Redshift doesn't support INDEX. pass + def column_sql(self, *args, **kwargs): + definition, params = super().column_sql(*args, **kwargs) + params = self._modify_params_for_redshift(params) + return definition, params + + def _modify_params_for_redshift(self, params): + """ + `Psycopg2.extensions.Binary(b'\x80\x00')` in params is converted to `'\\x80\\x00'::bytea` when applied to SQL placeholders. However, Redshift needs to treat binary columns as `to_varbyte('8000', 'hex')::varbyte` instead of `::bytea` [#]. + + [#]: https://docs.aws.amazon.com/redshift/latest/dg/r_VARBYTE_type.html + + So, this function converts `Binary` instances to `RedshiftBinary` instances. + RedshiftBinary is converted to `to_varbyte('8000', 'hex')::varbyte` when applied to placeholders. + """ + new_params = [ + RedshiftBinary(p.adapted) if isinstance(p, Binary) else p for p in params + ] + return new_params + def create_model(self, model): """ Takes a model and creates a table for it in the database. @@ -733,14 +758,31 @@ def _alter_column_with_recreate(self, model, old_field, new_field): }, params, ) + + type_cast = "" + if new_field.get_internal_type() == "BinaryField": + # In most cases, we don't change the type to a type that can't be cast, + # so we don't check it. + type_cast = "::" + _remove_length_from_type( + DatabaseWrapper.data_types["BinaryField"] + ) + elif ( + old_field.get_internal_type() == "BinaryField" + and new_field.get_internal_type() == "CharField" + ): + type_cast = "::" + _remove_length_from_type( + DatabaseWrapper.data_types["CharField"] + ) + # ## UPDATE SET 'tmp' = actions.append( ( - "UPDATE %(table)s SET %(new_column)s = %(old_column)s WHERE %(old_column)s IS NOT NULL" + "UPDATE %(table)s SET %(new_column)s = %(old_column)s%(cast)s WHERE %(old_column)s IS NOT NULL" % { "table": model._meta.db_table, "new_column": self.quote_name(new_field.column + "_tmp"), "old_column": self.quote_name(new_field.column), + "cast": type_cast, }, [], ) @@ -1080,6 +1122,7 @@ def _create_unique_sql( "BigAutoField": "bigint identity(1, 1)", "TextField": "varchar(max)", # text must be varchar(max) "UUIDField": "varchar(32)", # redshift doesn't support uuid fields + "BinaryField": "varbyte(%(max_length)s)", } diff --git a/django_redshift_backend/psycopg2adapter.py b/django_redshift_backend/psycopg2adapter.py new file mode 100644 index 0000000..4821f60 --- /dev/null +++ b/django_redshift_backend/psycopg2adapter.py @@ -0,0 +1,10 @@ +from codecs import encode + +from psycopg2.extensions import Binary + + +class RedshiftBinary(Binary): + def getquoted(self) -> bytes: + hex_encoded = encode(self.adapted, "hex_codec") + statement = b"to_varbyte('%s', 'hex')::varbyte" % hex_encoded + return statement diff --git a/examples/dj-sql-explorer/.env.sample b/examples/dj-sql-explorer/.env.sample index d89e73f..2d5a808 100644 --- a/examples/dj-sql-explorer/.env.sample +++ b/examples/dj-sql-explorer/.env.sample @@ -1,2 +1,3 @@ DATABASE_URL=redshift://user:password@...redshift.amazonaws.com:5439/?DISABLE_SERVER_SIDE_CURSORS=True SECRET_KEY=django-insecure-key +DEBUG=True diff --git a/examples/dj-sql-explorer/config/settings.py b/examples/dj-sql-explorer/config/settings.py index fd8ebd1..265cd82 100644 --- a/examples/dj-sql-explorer/config/settings.py +++ b/examples/dj-sql-explorer/config/settings.py @@ -139,3 +139,20 @@ # django-sql-explorer EXPLORER_CONNECTIONS = { 'Default': 'default' } EXPLORER_DEFAULT_CONNECTION = 'default' + +LOGGING = { + 'version': 1, + 'disable_existing_loggers': False, + 'handlers': { + 'console': { + 'level': 'DEBUG', + 'class': 'logging.StreamHandler', + }, + }, + 'loggers': { + 'django.db.backends': { + 'handlers': ['console'], + 'level': 'DEBUG', + }, + }, +} \ No newline at end of file diff --git a/examples/dj-sql-explorer/requirements.txt b/examples/dj-sql-explorer/requirements.txt index 4428a57..f29a551 100644 --- a/examples/dj-sql-explorer/requirements.txt +++ b/examples/dj-sql-explorer/requirements.txt @@ -1,4 +1,4 @@ -Django<4 +-e ../..[psycopg2-binary] django-environ==0.8.1 django-sql-explorer --e ../.. +python-dateutil>=2.9 diff --git a/examples/proj1/requirements.txt b/examples/proj1/requirements.txt index 78ffea4..ece3ff3 100644 --- a/examples/proj1/requirements.txt +++ b/examples/proj1/requirements.txt @@ -1,3 +1,2 @@ -Django<4 +-e ../..[psycopg2-binary] django-environ==0.8.1 --e ../.. diff --git a/tests/conftest.py b/tests/conftest.py index 7c6d614..49fca3c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -49,6 +49,10 @@ def postgres_fixture(): 'django_redshift_backend.base.DatabaseWrapper.data_types', BasePGDatabaseWrapper.data_types, ), \ + mock.patch( + 'django_redshift_backend.base.DatabaseSchemaEditor._modify_params_for_redshift', + lambda self, params: params + ), \ mock.patch( 'django_redshift_backend.base.DatabaseSchemaEditor._get_create_options', lambda self, model: '', diff --git a/tests/test_migrations.py b/tests/test_migrations.py index 0c34ece..62416fd 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -6,7 +6,7 @@ import pytest from test_base import OperationTestBase -from conftest import skipif_no_database, postgres_fixture +from conftest import skipif_no_database, postgres_fixture, TEST_WITH_POSTGRES, TEST_WITH_REDSHIFT @skipif_no_database @@ -289,6 +289,38 @@ def test_add_notnull_with_default(self): '''ALTER TABLE "test_pony" ADD COLUMN "name" varchar(10) DEFAULT '' NOT NULL;''', ], sqls) + @postgres_fixture() + def test_add_binary(self): + from django_redshift_backend.base import DatabaseWrapper, _remove_length_from_type + + new_state = self.set_up_test_model('test') + operations = [ + migrations.AddField( + model_name='Pony', + name='hash', + field=models.BinaryField( + max_length=10, + verbose_name='hash', + null=False, + default=b'\x80\x00', + ), + ), + ] + + with self.collect_sql() as sqls: + self.apply_operations('test', new_state, operations) + + bin_type = DatabaseWrapper.data_types['BinaryField'] % {"max_length": 10} + bin_cast = _remove_length_from_type(bin_type) + if TEST_WITH_POSTGRES: + default = fr"DEFAULT '\200\000'::{bin_cast}" + elif TEST_WITH_REDSHIFT: + default = fr"DEFAULT to_varbyte('8000', 'hex')::{bin_cast}" + + self.assertEqual([ + f'''ALTER TABLE "test_pony" ADD COLUMN "hash" {bin_type} {default} NOT NULL;''', + ], sqls) + @postgres_fixture() def test_alter_type(self): new_state = self.set_up_test_model('test') @@ -390,3 +422,85 @@ def test_alter_notnull_to_nullable(self): '''ALTER TABLE test_pony DROP COLUMN "weight" CASCADE;''', '''ALTER TABLE test_pony RENAME COLUMN "weight_tmp" TO "weight";''', ], sqls) + + @postgres_fixture() + def test_alter_type_char_to_binary(self): + from django_redshift_backend.base import DatabaseWrapper, _remove_length_from_type + + new_state = self.set_up_test_model('test') + operations = [ + migrations.AddField( + model_name='Pony', + name='hash', + field=models.CharField(max_length=10, verbose_name='hash', null=False, default=''), + ), + migrations.AlterField( + model_name='Pony', + name='hash', + field=models.BinaryField( + max_length=10, + verbose_name='hash', + null=False, + default=b'\x80\x00', + ), + ), + ] + + with self.collect_sql() as sqls: + self.apply_operations('test', new_state, operations) + + bin_type = DatabaseWrapper.data_types['BinaryField'] % {"max_length": 10} + bin_cast = _remove_length_from_type(bin_type) + if TEST_WITH_POSTGRES: + default = fr"DEFAULT '\200\000'::{bin_cast}" + elif TEST_WITH_REDSHIFT: + default = fr"DEFAULT to_varbyte('8000', 'hex')::{bin_cast}" + + self.assertEqual([ + '''ALTER TABLE "test_pony" ADD COLUMN "hash" varchar(10) DEFAULT '' NOT NULL;''', + f'''ALTER TABLE "test_pony" ADD COLUMN "hash_tmp" {bin_type} {default} NOT NULL;''', + f'''UPDATE test_pony SET "hash_tmp" = "hash"::{bin_cast} WHERE "hash" IS NOT NULL;''', + '''ALTER TABLE test_pony DROP COLUMN "hash" CASCADE;''', + '''ALTER TABLE test_pony RENAME COLUMN "hash_tmp" TO "hash";''', + ], sqls) + + @postgres_fixture() + def test_alter_type_binary_to_char(self): + from django_redshift_backend.base import DatabaseWrapper, _remove_length_from_type + + new_state = self.set_up_test_model('test') + operations = [ + migrations.AddField( + model_name='Pony', + name='hash', + field=models.BinaryField( + max_length=10, + verbose_name='hash', + null=False, + default=b'\x80\x00', + ), + ), + migrations.AlterField( + model_name='Pony', + name='hash', + field=models.CharField(max_length=10, verbose_name='hash', null=False, default=''), + ), + ] + + with self.collect_sql() as sqls: + self.apply_operations('test', new_state, operations) + + bin_type = DatabaseWrapper.data_types['BinaryField'] % {"max_length": 10} + bin_cast = _remove_length_from_type(bin_type) + if TEST_WITH_POSTGRES: + default = fr"DEFAULT '\200\000'::{bin_cast}" + elif TEST_WITH_REDSHIFT: + default = fr"DEFAULT to_varbyte('8000', 'hex')::{bin_cast}" + + self.assertEqual([ + f'''ALTER TABLE "test_pony" ADD COLUMN "hash" {bin_type} {default} NOT NULL;''', + '''ALTER TABLE "test_pony" ADD COLUMN "hash_tmp" varchar(10) DEFAULT '' NOT NULL;''', + '''UPDATE test_pony SET "hash_tmp" = "hash"::varchar WHERE "hash" IS NOT NULL;''', + '''ALTER TABLE test_pony DROP COLUMN "hash" CASCADE;''', + '''ALTER TABLE test_pony RENAME COLUMN "hash_tmp" TO "hash";''', + ], sqls)