#!/usr/bin/python
import os
import psycopg2


SCHEMA_DIRECTORY = '/usr/share/zabbix-postgresql'
SCHEMA_FILES = [
    'schema.sql',
]
DATA_FILES = [
    'images.sql',
    'data.sql',
]

def main():
    module = AnsibleModule(
        argument_spec=dict(
            database=dict(required=True),
            username=dict(required=True),
            password=dict(no_log=True),
            host=dict(),
            schema_only=dict(type='bool', default=False),
        ),
        supports_check_mode=True,
    )

    database = module.params['database']
    username = module.params['username']
    password = module.params['password']
    host = module.params['host']
    schema_only = module.params['schema_only']

    dsn = {
        'dbname': database,
        'user': username,
    }
    if password:
        dsn['password'] = password
    if host:
        dsn['host'] = host

    conn = psycopg2.connect(**dsn)
    with conn.cursor() as cur:
        cur.execute('SELECT tablename FROM pg_tables '
                    "WHERE schemaname = 'public'")
        changed = cur.rowcount == 0

    if changed and not module.check_mode:
        psql = module.get_bin_path('psql')
        if password:
            os.environ['PGPASSWORD'] = password
        if schema_only:
            scripts = SCHEMA_FILES
        else:
            scripts = SCHEMA_FILES + DATA_FILES
        for filename in scripts:
            path = os.path.join(SCHEMA_DIRECTORY, filename)
            cmd = ['psql', '-U', username, '-f', path]
            if host:
                cmd.extend(('-h', host))
            cmd.append(database)
            module.run_command(cmd, check_rc=True)

    module.exit_json(changed=changed)


from ansible.module_utils.basic import *

main()