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 

     8 def getstatusoutput(cmd):

     9     "A simpler version of getstatusoutput that works on win32 platforms."

    10     stdin, stdout, stderr = os.popen3(cmd)

    11     output = stdout.read()

    12     if output.endswith('\n'): output = output[:-1]

    13     status = stdin.close()

    14     return status, output

    15 

    16 def create_lang(db_name, verbosity=1):

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

    18 

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

    20     options = get_cmd_options(db_name)

    21 

    22     # Constructing the 'createlang' command.

    23     createlang_cmd = 'createlang %splpgsql' % options

    24     if verbosity >= 1: print createlang_cmd

    25 

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

    27     #  also be in your path.

    28     status, output = getstatusoutput(createlang_cmd)

    29 

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

    31     if status:

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

    33 

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

    35     "Creates database with psycopg2 cursor."

    36 

    37     # Constructing the necessary SQL to create the database (the DATABASE_USER

    38     #  must possess the privileges to create a database)

    39     create_sql = 'CREATE DATABASE %s' % connection.ops.quote_name(db_name)

    40     if settings.DATABASE_USER:

    41         create_sql += ' OWNER %s' % settings.DATABASE_USER

    42 

    43     cursor = connection.cursor()

    44     connection.creation.set_autocommit()

    45 

    46     try:

    47         # Trying to create the database first.

    48         cursor.execute(create_sql)

    49         #print create_sql

    50     except Exception, e:

    51         # Drop and recreate, if necessary.

    52         if not autoclobber:

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

    54         if autoclobber or confirm == 'yes':

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

    56             drop_db(db_name)

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

    58             cursor.execute(create_sql)

    59         else:

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

    61 

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

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

    64     """

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

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

    67      creation of the spatial database.

    68     """

    69 

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

    71     options = get_cmd_options(False)

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

    73     if verbosity >= 1: print create_cmd

    74 

    75     # Attempting to create the database.

    76     status, output = getstatusoutput(create_cmd)

    77 

    78     if status:

    79         if created_regex.match(output):

    80             if not autoclobber:

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

    82             if autoclobber or confirm == 'yes':

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

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

    85                 status, output = getstatusoutput(drop_cmd)

    86                 if status != 0:

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

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

    89                 status, output = getstatusoutput(create_cmd)

    90                 if status != 0:

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

    92             else:

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

    94         else:

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

    96 

    97 def create_spatial_db(test=False, verbosity=1, autoclobber=False, interactive=False):

    98     "Creates a spatial database based on the settings."

    99 

   100     # Making sure we're using PostgreSQL and psycopg2

   101     if settings.DATABASE_ENGINE != 'postgresql_psycopg2':

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

   103 

   104     # Getting the spatial database name

   105     if test:

   106         db_name = get_spatial_db(test=True)

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

   108     else:

   109         db_name = get_spatial_db()

   110         _create_with_shell(db_name, verbosity=verbosity, autoclobber=autoclobber)

   111 

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

   113     #  since the PostGIS installer enables this capability.

   114     if os.name != 'nt':

   115         create_lang(db_name, verbosity=verbosity)

   116 

   117     # Now adding in the PostGIS routines.

   118     load_postgis_sql(db_name, verbosity=verbosity)

   119 

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

   121 

   122     # Closing the connection

   123     connection.close()

   124     settings.DATABASE_NAME = db_name

   125 

   126     # Syncing the database

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

   128 

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

   130     """

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

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

   133     """

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

   135     cursor = connection.cursor()

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

   137 

   138 def get_cmd_options(db_name):

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

   140     # The db_name parameter is optional

   141     options = ''

   142     if db_name:

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

   144     if settings.DATABASE_USER:

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

   146     if settings.DATABASE_HOST:

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

   148     if settings.DATABASE_PORT:

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

   150     return options

   151 

   152 def get_spatial_db(test=False):

   153     """

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

   155      to return the test spatial database name.

   156     """

   157     if test:

   158         if settings.TEST_DATABASE_NAME:

   159             test_db_name = settings.TEST_DATABASE_NAME

   160         else:

   161             test_db_name = TEST_DATABASE_PREFIX + settings.DATABASE_NAME

   162         return test_db_name

   163     else:

   164         if not settings.DATABASE_NAME:

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

   166         return settings.DATABASE_NAME

   167 

   168 def load_postgis_sql(db_name, verbosity=1):

   169     """

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

   171      spatial_ref_sys.sql.

   172     """

   173 

   174     # Getting the path to the PostGIS SQL

   175     try:

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

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

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

   179         sql_path = settings.POSTGIS_SQL_PATH

   180     except AttributeError:

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

   182         if status:

   183             sql_path = '/usr/local/share'

   184 

   185     # The PostGIS SQL post-creation files.

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

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

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

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

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

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

   192 

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

   194     options = get_cmd_options(db_name)

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

   196 

   197     # Now trying to load up the PostGIS functions

   198     cmd = cmd_fmt % lwpostgis_file

   199     if verbosity >= 1: print cmd

   200     status, output = getstatusoutput(cmd)

   201     if status:

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

   203 

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

   205     cmd = cmd_fmt % srefsys_file

   206     if verbosity >= 1: print cmd

   207     status, output = getstatusoutput(cmd)

   208     if status:

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

   210 

   211     # Setting the permissions because on Windows platforms the owner

   212     #  of the spatial_ref_sys and geometry_columns tables is always

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

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

   215 

   216 def set_permissions(db_name):

   217     """

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

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

   220     """

   221     cursor = connection.cursor()

   222     user = settings.DATABASE_USER

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

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