1 import os, re, sys

     2 

     3 from django.conf import settings

     4 from django.core.management import call_command

     5 from django.db import connection

     6 from django.db.backends.creation import TEST_DATABASE_PREFIX

     7 from django.contrib.gis.db.backend.util import getstatusoutput

     8 

     9 def create_lang(db_name, verbosity=1):

    10     "Sets up the pl/pgsql language on the given database."

    11 

    12     # Getting the command-line options for the shell command

    13     options = get_cmd_options(db_name)

    14 

    15     # Constructing the 'createlang' command.

    16     createlang_cmd = 'createlang %splpgsql' % options

    17     if verbosity >= 1: print createlang_cmd

    18 

    19     # Must have database super-user privileges to execute createlang -- it must

    20     # also be in your path.

    21     status, output = getstatusoutput(createlang_cmd)

    22 

    23     # Checking the status of the command, 0 => execution successful

    24     if status:

    25         raise Exception("Error executing 'plpgsql' command: %s\n" % output)

    26 

    27 def _create_with_cursor(db_name, verbosity=1, autoclobber=False):

    28     "Creates database with psycopg2 cursor."

    29     qn = connection.ops.quote_name

    30 

    31     # Constructing the necessary SQL to create the database.

    32     create_sql = 'CREATE DATABASE %s' % qn(db_name)

    33 

    34     # If there's a template database for PostGIS set, then use it.

    35     if hasattr(settings, 'POSTGIS_TEMPLATE'):

    36         create_sql += ' TEMPLATE %s' % qn(settings.POSTGIS_TEMPLATE)

    37 

    38     # The DATABASE_USER must possess the privileges to create a spatial database.

    39     if settings.DATABASE_USER:

    40         create_sql += ' OWNER %s' % qn(settings.DATABASE_USER)

    41 

    42     cursor = connection.cursor()

    43     connection.creation.set_autocommit()

    44 

    45     try:

    46         # Trying to create the database first.

    47         cursor.execute(create_sql)

    48     except Exception, e:

    49         if 'already exists' in e.pgerror.lower():

    50             # Database already exists, drop and recreate if user agrees.

    51             if not autoclobber:

    52                 confirm = raw_input("\nIt appears the database, %s, already exists. Type 'yes' to delete it, or 'no' to cancel: " % db_name)

    53             if autoclobber or confirm == 'yes':

    54                 if verbosity >= 1: print 'Destroying old spatial database...'

    55                 drop_db(db_name)

    56                 if verbosity >= 1: print 'Creating new spatial database...'

    57                 cursor.execute(create_sql)

    58             else:

    59                 raise Exception('Spatial database creation canceled.')

    60         else:

    61             raise Exception('Spatial database creation failed: "%s"' % e.pgerror.strip())

    62 

    63 created_regex = re.compile(r'^createdb: database creation failed: ERROR:  database ".+" already exists')

    64 def _create_with_shell(db_name, verbosity=1, autoclobber=False):

    65     """

    66     If no spatial database already exists, then using a cursor will not work.

    67     Thus, a `createdb` command will be issued through the shell to bootstrap

    68     creation of the spatial database.

    69 

    70     TODO: Actually allow this method to be used without a spatial database

    71     in place first.

    72     """

    73     # Getting the command-line options for the shell command

    74     options = get_cmd_options(False)

    75     if hasattr(settings, 'POSTGIS_TEMPLATE'):

    76         options += '-T %s ' % settings.POSTGIS_TEMPlATE

    77 

    78     create_cmd = 'createdb -O %s %s%s' % (settings.DATABASE_USER, options, db_name)

    79     if verbosity >= 1: print create_cmd

    80 

    81     # Attempting to create the database.

    82     status, output = getstatusoutput(create_cmd)

    83 

    84     if status:

    85         if created_regex.match(output):

    86             if not autoclobber:

    87                 confirm = raw_input("\nIt appears the database, %s, already exists. Type 'yes' to delete it, or 'no' to cancel: " % db_name)

    88             if autoclobber or confirm == 'yes':

    89                 if verbosity >= 1: print 'Destroying old spatial database...'

    90                 drop_cmd = 'dropdb %s%s' % (options, db_name)

    91                 status, output = getstatusoutput(drop_cmd)

    92                 if status != 0:

    93                     raise Exception('Could not drop database %s: %s' % (db_name, output))

    94                 if verbosity >= 1: print 'Creating new spatial database...'

    95                 status, output = getstatusoutput(create_cmd)

    96                 if status != 0:

    97                     raise Exception('Could not create database after dropping: %s' % output)

    98             else:

    99                 raise Exception('Spatial Database Creation canceled.')

   100         else:

   101             raise Exception('Unknown error occurred in creating database: %s' % output)

   102 

   103 def create_test_spatial_db(verbosity=1, autoclobber=False, interactive=False):

   104     "Creates a test spatial database based on the settings."

   105 

   106     # Making sure we're using PostgreSQL and psycopg2

   107     if settings.DATABASE_ENGINE != 'postgresql_psycopg2':

   108         raise Exception('Spatial database creation only supported postgresql_psycopg2 platform.')

   109 

   110     # Getting the spatial database name

   111     db_name = get_spatial_db(test=True)

   112     _create_with_cursor(db_name, verbosity=verbosity, autoclobber=autoclobber)

   113 

   114     # If a template database is used, then don't need to do any of the following.

   115     if not hasattr(settings, 'POSTGIS_TEMPLATE'):

   116         # Creating the db language, does not need to be done on NT platforms

   117         # since the PostGIS installer enables this capability.

   118         if os.name != 'nt':

   119             create_lang(db_name, verbosity=verbosity)

   120 

   121         # Now adding in the PostGIS routines.

   122         load_postgis_sql(db_name, verbosity=verbosity)

   123 

   124     if verbosity >= 1: print 'Creation of spatial database %s successful.' % db_name

   125 

   126     # Closing the connection

   127     connection.close()

   128     settings.DATABASE_NAME = db_name

   129     connection.settings_dict["DATABASE_NAME"] = db_name

   130     can_rollback = connection.creation._rollback_works()

   131     settings.DATABASE_SUPPORTS_TRANSACTIONS = can_rollback

   132     connection.settings_dict["DATABASE_SUPPORTS_TRANSACTIONS"] = can_rollback

   133 

   134     # Syncing the database

   135     call_command('syncdb', verbosity=verbosity, interactive=interactive)

   136 

   137 def drop_db(db_name=False, test=False):

   138     """

   139     Drops the given database (defaults to what is returned from

   140     get_spatial_db()). All exceptions are propagated up to the caller.

   141     """

   142     if not db_name: db_name = get_spatial_db(test=test)

   143     cursor = connection.cursor()

   144     cursor.execute('DROP DATABASE %s' % connection.ops.quote_name(db_name))

   145 

   146 def get_cmd_options(db_name):

   147     "Obtains the command-line PostgreSQL connection options for shell commands."

   148     # The db_name parameter is optional

   149     options = ''

   150     if db_name:

   151         options += '-d %s ' % db_name

   152     if settings.DATABASE_USER:

   153         options += '-U %s ' % settings.DATABASE_USER

   154     if settings.DATABASE_HOST:

   155         options += '-h %s ' % settings.DATABASE_HOST

   156     if settings.DATABASE_PORT:

   157         options += '-p %s ' % settings.DATABASE_PORT

   158     return options

   159 

   160 def get_spatial_db(test=False):

   161     """

   162     Returns the name of the spatial database.  The 'test' keyword may be set

   163     to return the test spatial database name.

   164     """

   165     if test:

   166         if settings.TEST_DATABASE_NAME:

   167             test_db_name = settings.TEST_DATABASE_NAME

   168         else:

   169             test_db_name = TEST_DATABASE_PREFIX + settings.DATABASE_NAME

   170         return test_db_name

   171     else:

   172         if not settings.DATABASE_NAME:

   173             raise Exception('must configure DATABASE_NAME in settings.py')

   174         return settings.DATABASE_NAME

   175 

   176 def load_postgis_sql(db_name, verbosity=1):

   177     """

   178     This routine loads up the PostGIS SQL files lwpostgis.sql and

   179     spatial_ref_sys.sql.

   180     """

   181     # Getting the path to the PostGIS SQL

   182     try:

   183         # POSTGIS_SQL_PATH may be placed in settings to tell GeoDjango where the

   184         # PostGIS SQL files are located.  This is especially useful on Win32

   185         # platforms since the output of pg_config looks like "C:/PROGRA~1/..".

   186         sql_path = settings.POSTGIS_SQL_PATH

   187     except AttributeError:

   188         status, sql_path = getstatusoutput('pg_config --sharedir')

   189         if status:

   190             sql_path = '/usr/local/share'

   191 

   192     # The PostGIS SQL post-creation files.

   193     lwpostgis_file = os.path.join(sql_path, 'lwpostgis.sql')

   194     srefsys_file = os.path.join(sql_path, 'spatial_ref_sys.sql')

   195     if not os.path.isfile(lwpostgis_file):

   196         raise Exception('Could not find PostGIS function definitions in %s' % lwpostgis_file)

   197     if not os.path.isfile(srefsys_file):

   198         raise Exception('Could not find PostGIS spatial reference system definitions in %s' % srefsys_file)

   199 

   200     # Getting the psql command-line options, and command format.

   201     options = get_cmd_options(db_name)

   202     cmd_fmt = 'psql %s-f "%%s"' % options

   203 

   204     # Now trying to load up the PostGIS functions

   205     cmd = cmd_fmt % lwpostgis_file

   206     if verbosity >= 1: print cmd

   207     status, output = getstatusoutput(cmd)

   208     if status:

   209         raise Exception('Error in loading PostGIS lwgeometry routines.')

   210 

   211     # Now trying to load up the Spatial Reference System table

   212     cmd = cmd_fmt % srefsys_file

   213     if verbosity >= 1: print cmd

   214     status, output = getstatusoutput(cmd)

   215     if status:

   216         raise Exception('Error in loading PostGIS spatial_ref_sys table.')

   217 

   218     # Setting the permissions because on Windows platforms the owner

   219     # of the spatial_ref_sys and geometry_columns tables is always

   220     # the postgres user, regardless of how the db is created.

   221     if os.name == 'nt': set_permissions(db_name)

   222 

   223 def set_permissions(db_name):

   224     """

   225     Sets the permissions on the given database to that of the user specified

   226     in the settings.  Needed specifically for PostGIS on Win32 platforms.

   227     """

   228     cursor = connection.cursor()

   229     user = settings.DATABASE_USER

   230     cursor.execute('ALTER TABLE geometry_columns OWNER TO %s' % user)

   231     cursor.execute('ALTER TABLE spatial_ref_sys OWNER TO %s' % user)