django postgres database/user initializer

init_db.py

#!/usr/bin/env python
"""
This creates the django postgres user and database, should be refactored as a manage.py plugin.


(c) 2011 Barry Allard <[email protected]> 

MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

http://www.opensource.org/licenses/mit-license.php

"""

import settings
import os
import pwd
import subprocess
import sys

def better_call(argv, check=False):
    try:
        if check: 
            return subprocess.check_call(argv)
        else
            return subprocess.call(argv)
    except:
        # Fix: Ctrl-C while reading a password corrupting stty
        #subprocess.call('stty echo')
        subprocess.call('stty echo', shell=True)
        raise


def which(filename):
    """
    Naïeé implementation of 'which.'
         
    returns /path/to/filename if exists
            None              if filename does not exist

    >>> which('ls')
    /bin/ls
    >>> which('does_not_exist')
    >>> 
    """
    for directory in os.getenv("PATH").split(':'):
        test_filename = os.path.join(directory, filename)
        if os.path.exists(test_filename):
            return test_filename
    return None


def get_user():
    """Returns the name of the current user."""
    return pwd.getpwuid(os.getuid())[0]


def become_user(argv, become_user='root', check=True, sudo=None):
    """
    This will ensure we are running as the user specified.

    sudo   None   Autodetect sudo/su
           True   Use sudo
           False  Use su

    check  True   raise an Exception if command fails
           False  return False if command fails
   
    
    returns True  calling process is the correct user
            False calling process is NOT the correct user
    
    """
    def sudo_args():
        args = ['sudo', '-i']

        # Don't try to ask for a password if running non-interactively.        
        if not sys.stdout.isatty():
            args += ['-n']
 
        return args


    def can_sudo():
        return subprocess.call(sudo_args() + ['--', '/bin/true'])


    def can_su():
        return subprocess.call(['su', '-s', '/bin/true'])


    def sh_quote(s):
        """
        Escape single quotes of s and then single quote s
    
        a'b -> 'a'\\''b'

        """
        return "'" + s.replace("'", "'\\''") + "'"


    current_user = get_user()

    if current_user == become_user:
        # Already the correct user, continue.
        return True

    # As another user, path is different.
    # Turn relative paths into absolute.
    command_args = [os.path.abspath(argv[0])] + argv[1:] 

    quoted_command = " ".join(map(sh_quotes, command_args))
   
    # Try to use sudo if available
    if sudo == None and which('sudo') and can_sudo():
          sudo = True

    # Try to use su if available
    if sudo == None and sys.stdout.isatty() and which('su') and can_su():
          sudo = False

    if sudo == None:
          raise Exception('Neither su nor sudo are usable.')
    
    if sudo:
        args = sudo_args()

        # Become a user other than root
        if become_user != 'root':
           args = args + ['-u', become_user]

        args = args + ['--', quoted_command]

    # sudo == False : su instead
    else:
        # Become a user other than root, su to them
        if become_user != 'root':
            args = ['su', '-', become_user, '-c', quoted_command]

        # Not already root, wrap the su - <user> call to avoid a password prompt
        if current_user != 'root':
            quoted_command = " ".join(map(sh_quotes, args))
            args = ['su', '-', '-c', quoted_command]

    # Wrap everything in nohup to avoid early process demise.
    args = ['nohup'] + args + ['|', 'cat']

    # Run the program specified by args[0] and arguments args[1:]
    better_call(args, check)

    # Command is complete, stop any subsequent code.
    return False




def main(argv):

    if not get_user('postgres'):
        raise Exception("Rerun as user postgres")

    if settings.DATABASE_USER == '' or settings.DATABASE_USER == None:
        raise Exception("DATABASE_USER in settings.py needs to be set.")

    if settings.DATABASE_NAME == '' or settings.DATABASE_NAME == None:
        raise Exception("DATABASE_NAME in settings.py needs to be set.")


    extra_args = ['-e']

    if settings.DATABASE_HOST != '' and settings.DATABASE_HOST != None:
        extra_args = extra_args + ['-h', settings.DATABASE_HOST]

    if settings.DATABASE_PORT != '' and settings.DATABASE_PORT != None:
        extra_args = extra_args + ['-p', settings.DATABASE_PORT]

    if settings.DATABASE_ENGINE != 'postgresql_psycopg2':
        if settings.DATABASE_ENGINE != 'postgresql':
            raise Exception("DATABASE_ENGINE in settings.py needs\
 to be postgresql or postgresql_psycopg2")


    commands = [
        (   "Drop test database",
            which('dropdb'),
            extra_args + ['test_' + settings.DATABASE_NAME],
            False,
        ),
        (   "Drop production database",
            which('dropdb'),
            extra_args + [settings.DATABASE_NAME],
            False,
        ),
        (   "Drop database user",
            which('dropuser'),
            extra_args + [settings.DATABASE_USER],
            False,
        ),
        (   "Create database user", 
            which('createuser'),
            extra_args + ['-S', '-d', '-R', '-l', '-i', '-P', '-E', settings.DATABASE_USER],
        ),
 
        (   "Create production database",
            which('createdb'),
            extra_args + ['-O', settings.DATABASE_USER, '-E', 'UTF8', settings.DATABASE_NAME],
        ),
    ]


    for command in commands:
        if len(command) == 4:
            check_result = command[3]
        else:
            check_result = True 
        print str(command[0]) + " ..."
        better_call([command[1]] + command[2], check_result)



### Module
if __name__ == "__main__":
    if become_user(sys.argv, 'postgres'): 
        sys.exit(main(sys.argv))

    # Terminates everything
    sys.exit(1)