wip schema changes

Author: timgraham

django_mongodb/schema.py

@@ -3,6 +3,7 @@
 from pymongo import ASCENDING, DESCENDING
 from pymongo.operations import IndexModel
 
+from .fields import EmbeddedModelField
 from .query import wrap_database_errors
 from .utils import OperationCollector
 
@@ -27,17 +28,27 @@ def create_model(self, model):
             if field.remote_field.through._meta.auto_created:
                 self.create_model(field.remote_field.through)
 
-    def _create_model_indexes(self, model):
+    def _create_model_indexes(self, model, current_path=None, parent_model=None):
         """
         Create all indexes (field indexes & uniques, Meta.index_together,
         Meta.unique_together, Meta.constraints, Meta.indexes) for the model.
+
+        If this is a recursive call to due to an embedded model, `current_path`
+        tracks the path that must be prepended to the index's column, and
+        `parent_model` tracks the collection to add the index/constraint to.
         """
         if not model._meta.managed or model._meta.proxy or model._meta.swapped:
             return
         # Field indexes and uniques
         for field in model._meta.local_fields:
+            if isinstance(field, EmbeddedModelField):
+                new_path = ".".join((current_path, field.column)) if current_path else field.column
+                self._create_model_indexes(
+                    field.embedded_model, parent_model=parent_model or model, current_path=new_path
+                )
             if self._field_should_be_indexed(model, field):
-                self._add_field_index(model, field)
+                column_name = f"{current_path}.{field.column}" if current_path else None
+                self._add_field_index(parent_model or model, field, column_name=column_name)
             elif self._field_should_have_unique(field):
                 self._add_field_unique(model, field)
         # Meta.index_together (RemovedInDjango51Warning)
@@ -162,7 +173,7 @@ def alter_unique_together(self, model, old_unique_together, new_unique_together)
             constraint = UniqueConstraint(fields=field_names, name=name)
             self.add_constraint(model, constraint)
 
-    def add_index(self, model, index, field=None, unique=False):
+    def add_index(self, model, index, *, field=None, unique=False, column_name=None):
         if index.contains_expressions:
             return
         kwargs = {}
@@ -171,14 +182,15 @@ def add_index(self, model, index, field=None, unique=False):
             # Indexing on $type matches the value of most SQL databases by
             # allowing multiple null values for the unique constraint.
             if field:
-                filter_expression[field.column] = {"$type": field.db_type(self.connection)}
+                column = column_name or field.column
+                filter_expression[column] = {"$type": field.db_type(self.connection)}
             else:
                 for field_name, _ in index.fields_orders:
                     field_ = model._meta.get_field(field_name)
                     filter_expression[field_.column] = {"$type": field_.db_type(self.connection)}
             kwargs = {"partialFilterExpression": filter_expression, "unique": True}
         index_orders = (
-            [(field.column, ASCENDING)]
+            [(column_name or field.column, ASCENDING)]
             if field
             else [
                 # order is "" if ASCENDING or "DESC" if DESCENDING (see
@@ -196,11 +208,11 @@ def _add_composed_index(self, model, field_names):
         idx.set_name_with_model(model)
         self.add_index(model, idx)
 
-    def _add_field_index(self, model, field):
+    def _add_field_index(self, model, field, *, column_name=None):
         """Add an index on a field with db_index=True."""
-        index = Index(fields=[field.name])
-        index.name = self._create_index_name(model._meta.db_table, [field.column])
-        self.add_index(model, index, field=field)
+        index = Index(fields=[column_name or field.name])
+        index.name = self._create_index_name(model._meta.db_table, [column_name or field.column])
+        self.add_index(model, index, field=field, column_name=column_name)
 
     def remove_index(self, model, index):
         if index.contains_expressions:

tests/model_fields_/models.py

@@ -34,6 +34,7 @@ class EmbeddedModel(models.Model):
 class Address(models.Model):
     city = models.CharField(max_length=20)
     state = models.CharField(max_length=2)
+    zip_code = models.IntegerField(db_index=True)
 
 
 class Author(models.Model):

tests/schema_/__init__.py

@@ -0,0 +1,36 @@
+from django.apps.registry import Apps
+from django.db import models
+
+from django_mongodb.fields import EmbeddedModelField
+
+# Because we want to test creation and deletion of these as separate things,
+# these models are all inserted into a separate Apps so the main test
+# runner doesn't migrate them.
+
+new_apps = Apps()
+
+
+class Address(models.Model):
+    city = models.CharField(max_length=20)
+    state = models.CharField(max_length=2)
+    zip_code = models.IntegerField(db_index=True)
+
+    class Meta:
+        apps = new_apps
+
+
+class Author(models.Model):
+    name = models.CharField(max_length=10)
+    age = models.IntegerField(db_index=True)
+    address = EmbeddedModelField(Address)
+
+    class Meta:
+        apps = new_apps
+
+
+class Book(models.Model):
+    name = models.CharField(max_length=100)
+    author = EmbeddedModelField(Author)
+
+    class Meta:
+        apps = new_apps

tests/schema_/models.py

@@ -0,0 +1,255 @@
+import itertools
+
+from django.db import (
+    connection,
+)
+from django.test import (
+    TransactionTestCase,
+)
+
+from .models import Address, Author, Book, new_apps
+
+
+class SchemaTests(TransactionTestCase):
+    available_apps = []
+    models = [Address, Author, Book]
+
+    # Utility functions
+
+    def setUp(self):
+        # local_models should contain test dependent model classes that will be
+        # automatically removed from the app cache on test tear down.
+        self.local_models = []
+        # isolated_local_models contains models that are in test methods
+        # decorated with @isolate_apps.
+        self.isolated_local_models = []
+
+    def tearDown(self):
+        # Delete any tables made for our models
+        self.delete_tables()
+        new_apps.clear_cache()
+        for model in new_apps.get_models():
+            model._meta._expire_cache()
+        if "schema" in new_apps.all_models:
+            for model in self.local_models:
+                for many_to_many in model._meta.many_to_many:
+                    through = many_to_many.remote_field.through
+                    if through and through._meta.auto_created:
+                        del new_apps.all_models["schema"][through._meta.model_name]
+                del new_apps.all_models["schema"][model._meta.model_name]
+        if self.isolated_local_models:
+            with connection.schema_editor() as editor:
+                for model in self.isolated_local_models:
+                    editor.delete_model(model)
+
+    def delete_tables(self):
+        "Deletes all model tables for our models for a clean test environment"
+        converter = connection.introspection.identifier_converter
+        with connection.schema_editor() as editor:
+            connection.disable_constraint_checking()
+            table_names = connection.introspection.table_names()
+            if connection.features.ignores_table_name_case:
+                table_names = [table_name.lower() for table_name in table_names]
+            for model in itertools.chain(SchemaTests.models, self.local_models):
+                tbl = converter(model._meta.db_table)
+                if connection.features.ignores_table_name_case:
+                    tbl = tbl.lower()
+                if tbl in table_names:
+                    editor.delete_model(model)
+                    table_names.remove(tbl)
+            connection.enable_constraint_checking()
+
+    def column_classes(self, model):
+        with connection.cursor() as cursor:
+            columns = {
+                d[0]: (connection.introspection.get_field_type(d[1], d), d)
+                for d in connection.introspection.get_table_description(
+                    cursor,
+                    model._meta.db_table,
+                )
+            }
+        # SQLite has a different format for field_type
+        for name, (type, desc) in columns.items():
+            if isinstance(type, tuple):
+                columns[name] = (type[0], desc)
+        return columns
+
+    def get_primary_key(self, table):
+        with connection.cursor() as cursor:
+            return connection.introspection.get_primary_key_column(cursor, table)
+
+    def get_indexes(self, table):
+        """
+        Get the indexes on the table using a new cursor.
+        """
+        with connection.cursor() as cursor:
+            return [
+                c["columns"][0]
+                for c in connection.introspection.get_constraints(cursor, table).values()
+                if c["index"] and len(c["columns"]) == 1
+            ]
+
+    def get_uniques(self, table):
+        with connection.cursor() as cursor:
+            return [
+                c["columns"][0]
+                for c in connection.introspection.get_constraints(cursor, table).values()
+                if c["unique"] and len(c["columns"]) == 1
+            ]
+
+    def get_constraints(self, table):
+        """
+        Get the constraints on a table using a new cursor.
+        """
+        with connection.cursor() as cursor:
+            return connection.introspection.get_constraints(cursor, table)
+
+    def get_constraints_for_column(self, model, column_name):
+        constraints = self.get_constraints(model._meta.db_table)
+        constraints_for_column = []
+        for name, details in constraints.items():
+            if details["columns"] == [column_name]:
+                constraints_for_column.append(name)
+        return sorted(constraints_for_column)
+
+    def get_constraint_opclasses(self, constraint_name):
+        with connection.cursor() as cursor:
+            sql = """
+                SELECT opcname
+                FROM pg_opclass AS oc
+                JOIN pg_index as i on oc.oid = ANY(i.indclass)
+                JOIN pg_class as c on c.oid = i.indexrelid
+                WHERE c.relname = %s
+            """
+            cursor.execute(sql, [constraint_name])
+            return [row[0] for row in cursor.fetchall()]
+
+    def check_added_field_default(
+        self,
+        schema_editor,
+        model,
+        field,
+        field_name,
+        expected_default,
+        cast_function=None,
+    ):
+        schema_editor.add_field(model, field)
+        database_default = connection.database[model._meta.db_table].find_one().get(field_name)
+        # cursor.execute(
+        #     "SELECT {} FROM {};".format(field_name, model._meta.db_table)
+        # )
+        # database_default = cursor.fetchall()[0][0]
+        if cast_function and type(database_default) is not type(expected_default):
+            database_default = cast_function(database_default)
+        self.assertEqual(database_default, expected_default)
+
+    def get_constraints_count(self, table, column, fk_to):
+        """
+        Return a dict with keys 'fks', 'uniques, and 'indexes' indicating the
+        number of foreign keys, unique constraints, and indexes on
+        `table`.`column`. The `fk_to` argument is a 2-tuple specifying the
+        expected foreign key relationship's (table, column).
+        """
+        with connection.cursor() as cursor:
+            constraints = connection.introspection.get_constraints(cursor, table)
+        counts = {"fks": 0, "uniques": 0, "indexes": 0}
+        for c in constraints.values():
+            if c["columns"] == [column]:
+                if c["foreign_key"] == fk_to:
+                    counts["fks"] += 1
+                if c["unique"]:
+                    counts["uniques"] += 1
+                elif c["index"]:
+                    counts["indexes"] += 1
+        return counts
+
+    def get_column_collation(self, table, column):
+        with connection.cursor() as cursor:
+            return next(
+                f.collation
+                for f in connection.introspection.get_table_description(cursor, table)
+                if f.name == column
+            )
+
+    def get_column_comment(self, table, column):
+        with connection.cursor() as cursor:
+            return next(
+                f.comment
+                for f in connection.introspection.get_table_description(cursor, table)
+                if f.name == column
+            )
+
+    def get_table_comment(self, table):
+        with connection.cursor() as cursor:
+            return next(
+                t.comment
+                for t in connection.introspection.get_table_list(cursor)
+                if t.name == table
+            )
+
+    def assert_column_comment_not_exists(self, table, column):
+        with connection.cursor() as cursor:
+            columns = connection.introspection.get_table_description(cursor, table)
+        self.assertFalse(any(c.name == column and c.comment for c in columns))
+
+    def assertIndexOrder(self, table, index, order):
+        constraints = self.get_constraints(table)
+        self.assertIn(index, constraints)
+        index_orders = constraints[index]["orders"]
+        self.assertTrue(
+            all(val == expected for val, expected in zip(index_orders, order, strict=True))
+        )
+
+    def assertForeignKeyExists(self, model, column, expected_fk_table, field="id"):
+        """
+        Fail if the FK constraint on `model.Meta.db_table`.`column` to
+        `expected_fk_table`.id doesn't exist.
+        """
+        if not connection.features.can_introspect_foreign_keys:
+            return
+        constraints = self.get_constraints(model._meta.db_table)
+        constraint_fk = None
+        for details in constraints.values():
+            if details["columns"] == [column] and details["foreign_key"]:
+                constraint_fk = details["foreign_key"]
+                break
+        self.assertEqual(constraint_fk, (expected_fk_table, field))
+
+    def assertForeignKeyNotExists(self, model, column, expected_fk_table):
+        if not connection.features.can_introspect_foreign_keys:
+            return
+        with self.assertRaises(AssertionError):
+            self.assertForeignKeyExists(model, column, expected_fk_table)
+
+    def assertTableExists(self, model):
+        self.assertIn(model._meta.db_table, connection.introspection.table_names())
+
+    def assertTableNotExists(self, model):
+        self.assertNotIn(model._meta.db_table, connection.introspection.table_names())
+
+    # Tests
+    def test_embedded_index(self):
+        """db_index on an embedded model."""
+        with connection.schema_editor() as editor:
+            # Create the table
+            editor.create_model(Book)
+            # The table is there
+            self.assertTableExists(Book)
+            # Embedded indexes are created.
+            self.assertEqual(
+                self.get_constraints_for_column(Book, "author.age"),
+                ["schema__book_author.age_dc08100b"],
+            )
+            self.assertEqual(
+                self.get_constraints_for_column(Book, "author.address.zip_code"),
+                ["schema__book_author.address.zip_code_7b9a9307"],
+            )
+            # Clean up that table
+            editor.delete_model(Author)
+            # Indexes are gone.
+            self.assertEqual(
+                self.get_constraints_for_column(Author, "author.address.zip_code"),
+                [],
+            )
+        # The table is gone
+        self.assertTableNotExists(Author)