import subprocess
from kokki import *

def setup():
    postgresql_locale = env.system.locales[0]
    for l in env.system.locales:
        if 'utf8' in l.lower() or 'utf-8' in l.lower():
            postgresql_locale = l
            break
    env.set_attributes({
        'postgresql84.locale': postgresql_locale,
    })

setup()

def postgresql_user(username, exists=True, password=None, **privileges):
    '''
    Defines a PostgreSQL user to either exist or not. Takes an optional
    password (ignored if None) and a dictionary of boolean privileges.
    The available privileges are: superuser, createdb, inherit and login.
    '''
    if exists:
        # Create role update SQL
        defaults = {
                'superuser': False, 'createdb': False,
                'createrole': False, 'inherit': True,
                'login': True,
            }
        defaults.update(privileges)
        sql_privileges = ['%s%s' % ('' if defaults[p] else 'NO', p.upper()) for p in defaults.keys()]
        sql = "ALTER ROLE %s %s" % (username, ' '.join(sql_privileges))
        if password is not None:
            sql = "%s WITH PASSWORD '%s'" % (sql)
        # Run update if nescessary
        Execute('alter postgresql84 user %s' % username,
            command='sudo -u postgres psql -c "%s"' % sql)
    else:
        sql = 'DROP ROLE IF EXISTS %s' % username
        Execute('drop postgresql84 user %s' % username,
            command='sudo -u postgres psql -c "%s"' % sql)

def postgresql_db(name, owner='postgres', exists=True):
    '''
    Defines a PostgreSQL database to either exist or not. Takes an
    optional database owner.
    '''
    if exists:
        exists_cmd = "sudo -u postgres psql -f /dev/null %s" % name
        print 'exists_cmd', exists_cmd
        retval = subprocess.call(exists_cmd, shell=True)
        print 'retval', retval
        # Only create the database if it does not exist already
        if retval is not 0:
            Execute('creating postgresql84 database %s' % name,
                command='sudo -u postgres createdb %s' % name)
        # Update database owner
        Execute('updating postgresql84 database owner for %s' % name,
            command="sudo -u postgres psql -c 'ALTER DATABASE %s OWNER TO %s'" % (name, owner))
    else:
        Execute('dropping postgresql84 database "%s" if it exists' % name,
            command="sudo -u postgres psql -c 'DROP DATABASE IF EXISTS %s" % name)