Skip to content

Commit

Permalink
Merge pull request #95 from smzkw/dev/inspectdb-test
Browse files Browse the repository at this point in the history
add inspectdb test
  • Loading branch information
shimizukawa committed Feb 16, 2022
2 parents 58981ac + f33c35d commit 392a096
Show file tree
Hide file tree
Showing 5 changed files with 547 additions and 182 deletions.
23 changes: 15 additions & 8 deletions django_redshift_backend/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import uuid
import logging

import django
from django.conf import settings
from django.core.exceptions import FieldDoesNotExist
from django.db.backends.base.introspection import FieldInfo
Expand Down Expand Up @@ -594,10 +595,19 @@ def get_table_description(self, cursor, table_name):
AND n.nspname NOT IN ('pg_catalog', 'pg_toast')
AND pg_catalog.pg_table_is_visible(c.oid)
""", [table_name])
field_map = {
column_name: (is_nullable, column_default)
for (column_name, is_nullable, column_default) in cursor.fetchall()
}
# https://github.com/django/django/blob/stable/3.2.x/django/db/backends/postgresql/introspection.py#L85
# field_map = {line[0]: line[1:] for line in cursor.fetchall()}
field_map = {}
for column_name, is_nullable, column_default in cursor.fetchall():
_field_map = {
'null_ok': is_nullable,
'default': column_default,
}
if django.VERSION >= (3, 2):
# Redshift doesn't support user-defined collation
# https://docs.aws.amazon.com/redshift/latest/dg/c_collation_sequences.html
_field_map['collation'] = None
field_map[column_name] = _field_map
cursor.execute(
"SELECT * FROM %s LIMIT 1" % self.connection.ops.quote_name(table_name)
)
Expand All @@ -609,10 +619,7 @@ def get_table_description(self, cursor, table_name):
internal_size=column.internal_size,
precision=column.precision,
scale=column.scale,
null_ok=field_map[column.name][0],
default=field_map[column.name][1],
collation=None, # Redshift doesn't support user-defined collation
# https://docs.aws.amazon.com/redshift/latest/dg/c_collation_sequences.html
**field_map[column.name]
)
for column in cursor.description
]
Expand Down
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,6 @@ universal = 0
# 0 to make the generated wheels have `py3` tag

[flake8]
max-line-length=90
max-line-length=120
ignore = W504
exclude = tests/testapp/migrations
311 changes: 311 additions & 0 deletions tests/test_base.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,311 @@
# copy AS-IS from django repository at stable/3.2.x
# https://github.com/django/django/blob/754af45/tests/migrations/test_base.py
import os
import shutil
import tempfile
from contextlib import contextmanager
from importlib import import_module

from django.apps import apps
from django.db import connection, connections, migrations, models
from django.db.migrations.migration import Migration
from django.db.migrations.recorder import MigrationRecorder
from django.db.migrations.state import ProjectState
from django.test import TransactionTestCase
from django.test.utils import extend_sys_path
from django.utils.module_loading import module_dir


class MigrationTestBase(TransactionTestCase):
"""
Contains an extended set of asserts for testing migrations and schema operations.
"""

available_apps = ["migrations"]
databases = {'default', 'other'}

def tearDown(self):
# Reset applied-migrations state.
for db in self.databases:
recorder = MigrationRecorder(connections[db])
recorder.migration_qs.filter(app='migrations').delete()

def get_table_description(self, table, using='default'):
with connections[using].cursor() as cursor:
return connections[using].introspection.get_table_description(cursor, table)

def assertTableExists(self, table, using='default'):
with connections[using].cursor() as cursor:
self.assertIn(table, connections[using].introspection.table_names(cursor))

def assertTableNotExists(self, table, using='default'):
with connections[using].cursor() as cursor:
self.assertNotIn(table, connections[using].introspection.table_names(cursor))

def assertColumnExists(self, table, column, using='default'):
self.assertIn(column, [c.name for c in self.get_table_description(table, using=using)])

def assertColumnNotExists(self, table, column, using='default'):
self.assertNotIn(column, [c.name for c in self.get_table_description(table, using=using)])

def _get_column_allows_null(self, table, column, using):
return [c.null_ok for c in self.get_table_description(table, using=using) if c.name == column][0]

def assertColumnNull(self, table, column, using='default'):
self.assertTrue(self._get_column_allows_null(table, column, using))

def assertColumnNotNull(self, table, column, using='default'):
self.assertFalse(self._get_column_allows_null(table, column, using))

def assertIndexExists(self, table, columns, value=True, using='default', index_type=None):
with connections[using].cursor() as cursor:
self.assertEqual(
value,
any(
c["index"]
for c in connections[using].introspection.get_constraints(cursor, table).values()
if (
c['columns'] == list(columns) and
(index_type is None or c['type'] == index_type) and
not c['unique']
)
),
)

def assertIndexNotExists(self, table, columns):
return self.assertIndexExists(table, columns, False)

def assertIndexNameExists(self, table, index, using='default'):
with connections[using].cursor() as cursor:
self.assertIn(
index,
connection.introspection.get_constraints(cursor, table),
)

def assertIndexNameNotExists(self, table, index, using='default'):
with connections[using].cursor() as cursor:
self.assertNotIn(
index,
connection.introspection.get_constraints(cursor, table),
)

def assertConstraintExists(self, table, name, value=True, using='default'):
with connections[using].cursor() as cursor:
constraints = connections[using].introspection.get_constraints(cursor, table).items()
self.assertEqual(
value,
any(c['check'] for n, c in constraints if n == name),
)

def assertConstraintNotExists(self, table, name):
return self.assertConstraintExists(table, name, False)

def assertUniqueConstraintExists(self, table, columns, value=True, using='default'):
with connections[using].cursor() as cursor:
constraints = connections[using].introspection.get_constraints(cursor, table).values()
self.assertEqual(
value,
any(c['unique'] for c in constraints if c['columns'] == list(columns)),
)

def assertFKExists(self, table, columns, to, value=True, using='default'):
with connections[using].cursor() as cursor:
self.assertEqual(
value,
any(
c["foreign_key"] == to
for c in connections[using].introspection.get_constraints(cursor, table).values()
if c['columns'] == list(columns)
),
)

def assertFKNotExists(self, table, columns, to):
return self.assertFKExists(table, columns, to, False)

@contextmanager
def temporary_migration_module(self, app_label='migrations', module=None):
"""
Allows testing management commands in a temporary migrations module.
Wrap all invocations to makemigrations and squashmigrations with this
context manager in order to avoid creating migration files in your
source tree inadvertently.
Takes the application label that will be passed to makemigrations or
squashmigrations and the Python path to a migrations module.
The migrations module is used as a template for creating the temporary
migrations module. If it isn't provided, the application's migrations
module is used, if it exists.
Returns the filesystem path to the temporary migrations module.
"""
with tempfile.TemporaryDirectory() as temp_dir:
target_dir = tempfile.mkdtemp(dir=temp_dir)
with open(os.path.join(target_dir, '__init__.py'), 'w'):
pass
target_migrations_dir = os.path.join(target_dir, 'migrations')

if module is None:
module = apps.get_app_config(app_label).name + '.migrations'

try:
source_migrations_dir = module_dir(import_module(module))
except (ImportError, ValueError):
pass
else:
shutil.copytree(source_migrations_dir, target_migrations_dir)

with extend_sys_path(temp_dir):
new_module = os.path.basename(target_dir) + '.migrations'
with self.settings(MIGRATION_MODULES={app_label: new_module}):
yield target_migrations_dir


class OperationTestBase(MigrationTestBase):
"""Common functions to help test operations."""

@classmethod
def setUpClass(cls):
super().setUpClass()
cls._initial_table_names = frozenset(connection.introspection.table_names())

def tearDown(self):
self.cleanup_test_tables()
super().tearDown()

def cleanup_test_tables(self):
table_names = frozenset(connection.introspection.table_names()) - self._initial_table_names
with connection.schema_editor() as editor:
with connection.constraint_checks_disabled():
for table_name in table_names:
editor.execute(editor.sql_delete_table % {
'table': editor.quote_name(table_name),
})

def apply_operations(self, app_label, project_state, operations, atomic=True):
migration = Migration('name', app_label)
migration.operations = operations
with connection.schema_editor(atomic=atomic) as editor:
return migration.apply(project_state, editor)

def unapply_operations(self, app_label, project_state, operations, atomic=True):
migration = Migration('name', app_label)
migration.operations = operations
with connection.schema_editor(atomic=atomic) as editor:
return migration.unapply(project_state, editor)

def make_test_state(self, app_label, operation, **kwargs):
"""
Makes a test state using set_up_test_model and returns the
original state and the state after the migration is applied.
"""
project_state = self.set_up_test_model(app_label, **kwargs)
new_state = project_state.clone()
operation.state_forwards(app_label, new_state)
return project_state, new_state

def set_up_test_model(
self, app_label, second_model=False, third_model=False, index=False,
multicol_index=False, related_model=False, mti_model=False,
proxy_model=False, manager_model=False, unique_together=False,
options=False, db_table=None, index_together=False, constraints=None,
indexes=None,
):
"""Creates a test model state and database table."""
# Make the "current" state.
model_options = {
'swappable': 'TEST_SWAP_MODEL',
'index_together': [['weight', 'pink']] if index_together else [],
'unique_together': [['pink', 'weight']] if unique_together else [],
}
if options:
model_options['permissions'] = [('can_groom', 'Can groom')]
if db_table:
model_options['db_table'] = db_table
operations = [migrations.CreateModel(
'Pony',
[
('id', models.AutoField(primary_key=True)),
('pink', models.IntegerField(default=3)),
('weight', models.FloatField()),
],
options=model_options,
)]
if index:
operations.append(migrations.AddIndex(
'Pony',
models.Index(fields=['pink'], name='pony_pink_idx'),
))
if multicol_index:
operations.append(migrations.AddIndex(
'Pony',
models.Index(fields=['pink', 'weight'], name='pony_test_idx'),
))
if indexes:
for index in indexes:
operations.append(migrations.AddIndex('Pony', index))
if constraints:
for constraint in constraints:
operations.append(migrations.AddConstraint('Pony', constraint))
if second_model:
operations.append(migrations.CreateModel(
'Stable',
[
('id', models.AutoField(primary_key=True)),
]
))
if third_model:
operations.append(migrations.CreateModel(
'Van',
[
('id', models.AutoField(primary_key=True)),
]
))
if related_model:
operations.append(migrations.CreateModel(
'Rider',
[
('id', models.AutoField(primary_key=True)),
('pony', models.ForeignKey('Pony', models.CASCADE)),
('friend', models.ForeignKey('self', models.CASCADE, null=True))
],
))
if mti_model:
operations.append(migrations.CreateModel(
'ShetlandPony',
fields=[
('pony_ptr', models.OneToOneField(
'Pony',
models.CASCADE,
auto_created=True,
parent_link=True,
primary_key=True,
to_field='id',
serialize=False,
)),
('cuteness', models.IntegerField(default=1)),
],
bases=['%s.Pony' % app_label],
))
if proxy_model:
operations.append(migrations.CreateModel(
'ProxyPony',
fields=[],
options={'proxy': True},
bases=['%s.Pony' % app_label],
))
if manager_model:
from .models import FoodManager, FoodQuerySet
operations.append(migrations.CreateModel(
'Food',
fields=[
('id', models.AutoField(primary_key=True)),
],
managers=[
('food_qs', FoodQuerySet.as_manager()),
('food_mgr', FoodManager('a', 'b')),
('food_mgr_kwargs', FoodManager('x', 'y', 3, 4)),
]
))
return self.apply_operations(app_label, ProjectState(), operations)

0 comments on commit 392a096

Please sign in to comment.