creation.py 4.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. from django.conf import settings
  2. from django.db.backends.postgresql_psycopg2.creation import DatabaseCreation
  3. from django.utils.functional import cached_property
  4. class PostGISCreation(DatabaseCreation):
  5. geom_index_type = 'GIST'
  6. geom_index_ops = 'GIST_GEOMETRY_OPS'
  7. geom_index_ops_nd = 'GIST_GEOMETRY_OPS_ND'
  8. @cached_property
  9. def template_postgis(self):
  10. template_postgis = getattr(settings, 'POSTGIS_TEMPLATE', 'template_postgis')
  11. with self.connection.cursor() as cursor:
  12. cursor.execute('SELECT 1 FROM pg_database WHERE datname = %s LIMIT 1;', (template_postgis,))
  13. if cursor.fetchone():
  14. return template_postgis
  15. return None
  16. def sql_indexes_for_field(self, model, f, style):
  17. "Return any spatial index creation SQL for the field."
  18. from django.contrib.gis.db.models.fields import GeometryField
  19. output = super(PostGISCreation, self).sql_indexes_for_field(model, f, style)
  20. if isinstance(f, GeometryField):
  21. gqn = self.connection.ops.geo_quote_name
  22. qn = self.connection.ops.quote_name
  23. db_table = model._meta.db_table
  24. if f.geography or self.connection.ops.geometry:
  25. # Geography and Geometry (PostGIS 2.0+) columns are
  26. # created normally.
  27. pass
  28. else:
  29. # Geometry columns are created by `AddGeometryColumn`
  30. # stored procedure.
  31. output.append(style.SQL_KEYWORD('SELECT ') +
  32. style.SQL_TABLE('AddGeometryColumn') + '(' +
  33. style.SQL_TABLE(gqn(db_table)) + ', ' +
  34. style.SQL_FIELD(gqn(f.column)) + ', ' +
  35. style.SQL_FIELD(str(f.srid)) + ', ' +
  36. style.SQL_COLTYPE(gqn(f.geom_type)) + ', ' +
  37. style.SQL_KEYWORD(str(f.dim)) + ');')
  38. if not f.null:
  39. # Add a NOT NULL constraint to the field
  40. output.append(style.SQL_KEYWORD('ALTER TABLE ') +
  41. style.SQL_TABLE(qn(db_table)) +
  42. style.SQL_KEYWORD(' ALTER ') +
  43. style.SQL_FIELD(qn(f.column)) +
  44. style.SQL_KEYWORD(' SET NOT NULL') + ';')
  45. if f.spatial_index:
  46. # Spatial indexes created the same way for both Geometry and
  47. # Geography columns.
  48. # PostGIS 2.0 does not support GIST_GEOMETRY_OPS. So, on 1.5
  49. # we use GIST_GEOMETRY_OPS, on 2.0 we use either "nd" ops
  50. # which are fast on multidimensional cases, or just plain
  51. # gist index for the 2d case.
  52. if f.geography:
  53. index_ops = ''
  54. elif self.connection.ops.geometry:
  55. if f.dim > 2:
  56. index_ops = ' ' + style.SQL_KEYWORD(self.geom_index_ops_nd)
  57. else:
  58. index_ops = ''
  59. else:
  60. index_ops = ' ' + style.SQL_KEYWORD(self.geom_index_ops)
  61. output.append(style.SQL_KEYWORD('CREATE INDEX ') +
  62. style.SQL_TABLE(qn('%s_%s_id' % (db_table, f.column))) +
  63. style.SQL_KEYWORD(' ON ') +
  64. style.SQL_TABLE(qn(db_table)) +
  65. style.SQL_KEYWORD(' USING ') +
  66. style.SQL_COLTYPE(self.geom_index_type) + ' ( ' +
  67. style.SQL_FIELD(qn(f.column)) + index_ops + ' );')
  68. return output
  69. def sql_table_creation_suffix(self):
  70. if self.template_postgis is not None:
  71. return ' TEMPLATE %s' % (
  72. self.connection.ops.quote_name(self.template_postgis),)
  73. return ''
  74. def _create_test_db(self, verbosity, autoclobber, keepdb=False):
  75. test_database_name = super(PostGISCreation, self)._create_test_db(verbosity, autoclobber, keepdb)
  76. if keepdb:
  77. return test_database_name
  78. if self.template_postgis is None:
  79. # Connect to the test database in order to create the postgis extension
  80. self.connection.close()
  81. self.connection.settings_dict["NAME"] = test_database_name
  82. with self.connection.cursor() as cursor:
  83. cursor.execute("CREATE EXTENSION IF NOT EXISTS postgis")
  84. cursor.connection.commit()
  85. return test_database_name