1. from unittest import skipUnless
    
  2. 
    
  3. from django.contrib.gis.db.models import fields
    
  4. from django.contrib.gis.geos import MultiPolygon, Polygon
    
  5. from django.core.exceptions import ImproperlyConfigured
    
  6. from django.db import connection, migrations, models
    
  7. from django.db.migrations.migration import Migration
    
  8. from django.db.migrations.state import ProjectState
    
  9. from django.test import TransactionTestCase, skipIfDBFeature, skipUnlessDBFeature
    
  10. 
    
  11. try:
    
  12.     GeometryColumns = connection.ops.geometry_columns()
    
  13.     HAS_GEOMETRY_COLUMNS = True
    
  14. except NotImplementedError:
    
  15.     HAS_GEOMETRY_COLUMNS = False
    
  16. 
    
  17. 
    
  18. class OperationTestCase(TransactionTestCase):
    
  19.     available_apps = ["gis_tests.gis_migrations"]
    
  20.     get_opclass_query = """
    
  21.         SELECT opcname, c.relname FROM pg_opclass AS oc
    
  22.         JOIN pg_index as i on oc.oid = ANY(i.indclass)
    
  23.         JOIN pg_class as c on c.oid = i.indexrelid
    
  24.         WHERE c.relname = %s
    
  25.     """
    
  26. 
    
  27.     def tearDown(self):
    
  28.         # Delete table after testing
    
  29.         if hasattr(self, "current_state"):
    
  30.             self.apply_operations(
    
  31.                 "gis", self.current_state, [migrations.DeleteModel("Neighborhood")]
    
  32.             )
    
  33.         super().tearDown()
    
  34. 
    
  35.     @property
    
  36.     def has_spatial_indexes(self):
    
  37.         if connection.ops.mysql:
    
  38.             with connection.cursor() as cursor:
    
  39.                 return connection.introspection.supports_spatial_index(
    
  40.                     cursor, "gis_neighborhood"
    
  41.                 )
    
  42.         return True
    
  43. 
    
  44.     def get_table_description(self, table):
    
  45.         with connection.cursor() as cursor:
    
  46.             return connection.introspection.get_table_description(cursor, table)
    
  47. 
    
  48.     def assertColumnExists(self, table, column):
    
  49.         self.assertIn(column, [c.name for c in self.get_table_description(table)])
    
  50. 
    
  51.     def assertColumnNotExists(self, table, column):
    
  52.         self.assertNotIn(column, [c.name for c in self.get_table_description(table)])
    
  53. 
    
  54.     def apply_operations(self, app_label, project_state, operations):
    
  55.         migration = Migration("name", app_label)
    
  56.         migration.operations = operations
    
  57.         with connection.schema_editor() as editor:
    
  58.             return migration.apply(project_state, editor)
    
  59. 
    
  60.     def set_up_test_model(self, force_raster_creation=False):
    
  61.         test_fields = [
    
  62.             ("id", models.AutoField(primary_key=True)),
    
  63.             ("name", models.CharField(max_length=100, unique=True)),
    
  64.             ("geom", fields.MultiPolygonField(srid=4326)),
    
  65.         ]
    
  66.         if connection.features.supports_raster or force_raster_creation:
    
  67.             test_fields += [("rast", fields.RasterField(srid=4326, null=True))]
    
  68.         operations = [migrations.CreateModel("Neighborhood", test_fields)]
    
  69.         self.current_state = self.apply_operations("gis", ProjectState(), operations)
    
  70. 
    
  71.     def assertGeometryColumnsCount(self, expected_count):
    
  72.         self.assertEqual(
    
  73.             GeometryColumns.objects.filter(
    
  74.                 **{
    
  75.                     "%s__iexact" % GeometryColumns.table_name_col(): "gis_neighborhood",
    
  76.                 }
    
  77.             ).count(),
    
  78.             expected_count,
    
  79.         )
    
  80. 
    
  81.     def assertSpatialIndexExists(self, table, column, raster=False):
    
  82.         with connection.cursor() as cursor:
    
  83.             constraints = connection.introspection.get_constraints(cursor, table)
    
  84.         if raster:
    
  85.             self.assertTrue(
    
  86.                 any(
    
  87.                     "st_convexhull(%s)" % column in c["definition"]
    
  88.                     for c in constraints.values()
    
  89.                     if c["definition"] is not None
    
  90.                 )
    
  91.             )
    
  92.         else:
    
  93.             self.assertIn([column], [c["columns"] for c in constraints.values()])
    
  94. 
    
  95.     def alter_gis_model(
    
  96.         self,
    
  97.         migration_class,
    
  98.         model_name,
    
  99.         field_name,
    
  100.         blank=False,
    
  101.         field_class=None,
    
  102.         field_class_kwargs=None,
    
  103.     ):
    
  104.         args = [model_name, field_name]
    
  105.         if field_class:
    
  106.             field_class_kwargs = field_class_kwargs or {"srid": 4326, "blank": blank}
    
  107.             args.append(field_class(**field_class_kwargs))
    
  108.         operation = migration_class(*args)
    
  109.         old_state = self.current_state.clone()
    
  110.         operation.state_forwards("gis", self.current_state)
    
  111.         with connection.schema_editor() as editor:
    
  112.             operation.database_forwards("gis", editor, old_state, self.current_state)
    
  113. 
    
  114. 
    
  115. class OperationTests(OperationTestCase):
    
  116.     def setUp(self):
    
  117.         super().setUp()
    
  118.         self.set_up_test_model()
    
  119. 
    
  120.     def test_add_geom_field(self):
    
  121.         """
    
  122.         Test the AddField operation with a geometry-enabled column.
    
  123.         """
    
  124.         self.alter_gis_model(
    
  125.             migrations.AddField, "Neighborhood", "path", False, fields.LineStringField
    
  126.         )
    
  127.         self.assertColumnExists("gis_neighborhood", "path")
    
  128. 
    
  129.         # Test GeometryColumns when available
    
  130.         if HAS_GEOMETRY_COLUMNS:
    
  131.             self.assertGeometryColumnsCount(2)
    
  132. 
    
  133.         # Test spatial indices when available
    
  134.         if self.has_spatial_indexes:
    
  135.             self.assertSpatialIndexExists("gis_neighborhood", "path")
    
  136. 
    
  137.     @skipUnless(HAS_GEOMETRY_COLUMNS, "Backend doesn't support GeometryColumns.")
    
  138.     def test_geom_col_name(self):
    
  139.         self.assertEqual(
    
  140.             GeometryColumns.geom_col_name(),
    
  141.             "column_name" if connection.ops.oracle else "f_geometry_column",
    
  142.         )
    
  143. 
    
  144.     @skipUnlessDBFeature("supports_raster")
    
  145.     def test_add_raster_field(self):
    
  146.         """
    
  147.         Test the AddField operation with a raster-enabled column.
    
  148.         """
    
  149.         self.alter_gis_model(
    
  150.             migrations.AddField, "Neighborhood", "heatmap", False, fields.RasterField
    
  151.         )
    
  152.         self.assertColumnExists("gis_neighborhood", "heatmap")
    
  153. 
    
  154.         # Test spatial indices when available
    
  155.         if self.has_spatial_indexes:
    
  156.             self.assertSpatialIndexExists("gis_neighborhood", "heatmap", raster=True)
    
  157. 
    
  158.     def test_add_blank_geom_field(self):
    
  159.         """
    
  160.         Should be able to add a GeometryField with blank=True.
    
  161.         """
    
  162.         self.alter_gis_model(
    
  163.             migrations.AddField, "Neighborhood", "path", True, fields.LineStringField
    
  164.         )
    
  165.         self.assertColumnExists("gis_neighborhood", "path")
    
  166. 
    
  167.         # Test GeometryColumns when available
    
  168.         if HAS_GEOMETRY_COLUMNS:
    
  169.             self.assertGeometryColumnsCount(2)
    
  170. 
    
  171.         # Test spatial indices when available
    
  172.         if self.has_spatial_indexes:
    
  173.             self.assertSpatialIndexExists("gis_neighborhood", "path")
    
  174. 
    
  175.     @skipUnlessDBFeature("supports_raster")
    
  176.     def test_add_blank_raster_field(self):
    
  177.         """
    
  178.         Should be able to add a RasterField with blank=True.
    
  179.         """
    
  180.         self.alter_gis_model(
    
  181.             migrations.AddField, "Neighborhood", "heatmap", True, fields.RasterField
    
  182.         )
    
  183.         self.assertColumnExists("gis_neighborhood", "heatmap")
    
  184. 
    
  185.         # Test spatial indices when available
    
  186.         if self.has_spatial_indexes:
    
  187.             self.assertSpatialIndexExists("gis_neighborhood", "heatmap", raster=True)
    
  188. 
    
  189.     def test_remove_geom_field(self):
    
  190.         """
    
  191.         Test the RemoveField operation with a geometry-enabled column.
    
  192.         """
    
  193.         self.alter_gis_model(migrations.RemoveField, "Neighborhood", "geom")
    
  194.         self.assertColumnNotExists("gis_neighborhood", "geom")
    
  195. 
    
  196.         # Test GeometryColumns when available
    
  197.         if HAS_GEOMETRY_COLUMNS:
    
  198.             self.assertGeometryColumnsCount(0)
    
  199. 
    
  200.     @skipUnlessDBFeature("supports_raster")
    
  201.     def test_remove_raster_field(self):
    
  202.         """
    
  203.         Test the RemoveField operation with a raster-enabled column.
    
  204.         """
    
  205.         self.alter_gis_model(migrations.RemoveField, "Neighborhood", "rast")
    
  206.         self.assertColumnNotExists("gis_neighborhood", "rast")
    
  207. 
    
  208.     def test_create_model_spatial_index(self):
    
  209.         if not self.has_spatial_indexes:
    
  210.             self.skipTest("No support for Spatial indexes")
    
  211. 
    
  212.         self.assertSpatialIndexExists("gis_neighborhood", "geom")
    
  213. 
    
  214.         if connection.features.supports_raster:
    
  215.             self.assertSpatialIndexExists("gis_neighborhood", "rast", raster=True)
    
  216. 
    
  217.     @skipUnlessDBFeature("supports_3d_storage")
    
  218.     def test_add_3d_field_opclass(self):
    
  219.         if not connection.ops.postgis:
    
  220.             self.skipTest("PostGIS-specific test.")
    
  221. 
    
  222.         self.alter_gis_model(
    
  223.             migrations.AddField,
    
  224.             "Neighborhood",
    
  225.             "point3d",
    
  226.             field_class=fields.PointField,
    
  227.             field_class_kwargs={"dim": 3},
    
  228.         )
    
  229.         self.assertColumnExists("gis_neighborhood", "point3d")
    
  230.         self.assertSpatialIndexExists("gis_neighborhood", "point3d")
    
  231. 
    
  232.         with connection.cursor() as cursor:
    
  233.             index_name = "gis_neighborhood_point3d_113bc868_id"
    
  234.             cursor.execute(self.get_opclass_query, [index_name])
    
  235.             self.assertEqual(
    
  236.                 cursor.fetchall(),
    
  237.                 [("gist_geometry_ops_nd", index_name)],
    
  238.             )
    
  239. 
    
  240.     @skipUnlessDBFeature("can_alter_geometry_field", "supports_3d_storage")
    
  241.     def test_alter_geom_field_dim(self):
    
  242.         Neighborhood = self.current_state.apps.get_model("gis", "Neighborhood")
    
  243.         p1 = Polygon(((0, 0), (0, 1), (1, 1), (1, 0), (0, 0)))
    
  244.         Neighborhood.objects.create(name="TestDim", geom=MultiPolygon(p1, p1))
    
  245.         # Add 3rd dimension.
    
  246.         self.alter_gis_model(
    
  247.             migrations.AlterField,
    
  248.             "Neighborhood",
    
  249.             "geom",
    
  250.             False,
    
  251.             fields.MultiPolygonField,
    
  252.             field_class_kwargs={"srid": 4326, "dim": 3},
    
  253.         )
    
  254.         self.assertTrue(Neighborhood.objects.first().geom.hasz)
    
  255.         # Rewind to 2 dimensions.
    
  256.         self.alter_gis_model(
    
  257.             migrations.AlterField,
    
  258.             "Neighborhood",
    
  259.             "geom",
    
  260.             False,
    
  261.             fields.MultiPolygonField,
    
  262.             field_class_kwargs={"srid": 4326, "dim": 2},
    
  263.         )
    
  264.         self.assertFalse(Neighborhood.objects.first().geom.hasz)
    
  265. 
    
  266.     @skipUnlessDBFeature(
    
  267.         "supports_column_check_constraints", "can_introspect_check_constraints"
    
  268.     )
    
  269.     def test_add_check_constraint(self):
    
  270.         Neighborhood = self.current_state.apps.get_model("gis", "Neighborhood")
    
  271.         poly = Polygon(((0, 0), (0, 1), (1, 1), (1, 0), (0, 0)))
    
  272.         constraint = models.CheckConstraint(
    
  273.             check=models.Q(geom=poly),
    
  274.             name="geom_within_constraint",
    
  275.         )
    
  276.         Neighborhood._meta.constraints = [constraint]
    
  277.         with connection.schema_editor() as editor:
    
  278.             editor.add_constraint(Neighborhood, constraint)
    
  279.         with connection.cursor() as cursor:
    
  280.             constraints = connection.introspection.get_constraints(
    
  281.                 cursor,
    
  282.                 Neighborhood._meta.db_table,
    
  283.             )
    
  284.             self.assertIn("geom_within_constraint", constraints)
    
  285. 
    
  286. 
    
  287. @skipIfDBFeature("supports_raster")
    
  288. class NoRasterSupportTests(OperationTestCase):
    
  289.     def test_create_raster_model_on_db_without_raster_support(self):
    
  290.         msg = "Raster fields require backends with raster support."
    
  291.         with self.assertRaisesMessage(ImproperlyConfigured, msg):
    
  292.             self.set_up_test_model(force_raster_creation=True)
    
  293. 
    
  294.     def test_add_raster_field_on_db_without_raster_support(self):
    
  295.         msg = "Raster fields require backends with raster support."
    
  296.         with self.assertRaisesMessage(ImproperlyConfigured, msg):
    
  297.             self.set_up_test_model()
    
  298.             self.alter_gis_model(
    
  299.                 migrations.AddField,
    
  300.                 "Neighborhood",
    
  301.                 "heatmap",
    
  302.                 False,
    
  303.                 fields.RasterField,
    
  304.             )