← Back to branch summary

~stub/charms/precise/postgresql/repoless

« back to all changes in this revision

Viewing changes to lib/test-client-charm/hooks/hooks.py

  • Committer: Stuart Bishop
  • Date: 2014-10-14 08:45:43 UTC
  • Revision ID: stuart@stuartbishop.net-20141014084543-vfalg205cyhho6bv
Embed test client charm and remove requirement of JUJU_REPOSITORY being set

Show diffs side-by-side

added added

removed removed

Lines of Context:
lib/test-client-charm/hooks/hooks.py
 
1
#!/usr/bin/env python
 
2
 
 
3
import os.path
 
4
import re
 
5
import shutil
 
6
import sys
 
7
from textwrap import dedent
 
8
 
 
9
from charmhelpers.core import hookenv, host
 
10
from charmhelpers.core.hookenv import log, DEBUG, INFO
 
11
from charmhelpers import fetch
 
12
 
 
13
 
 
14
CLIENT_RELATION_TYPES = frozenset(['db', 'db-admin'])
 
15
 
 
16
DATA_DIR = os.path.join(
 
17
    '/var/lib/units', hookenv.local_unit().replace('/', '-'))
 
18
SCRIPT_DIR = os.path.join(DATA_DIR, 'bin')
 
19
PGPASS_DIR = os.path.join(DATA_DIR, 'pgpass')
 
20
 
 
21
 
 
22
def update_system_path():
 
23
    org_lines = open('/etc/environment', 'rb').readlines()
 
24
    env_lines = []
 
25
 
 
26
    for line in org_lines:
 
27
        if line.startswith('PATH=') and SCRIPT_DIR not in line:
 
28
            line = re.sub(
 
29
                """(['"]?)$""",
 
30
                ":{}\\1".format(SCRIPT_DIR),
 
31
                line, 1)
 
32
        env_lines.append(line)
 
33
 
 
34
    if org_lines != env_lines:
 
35
        content = '\n'.join(env_lines)
 
36
        host.write_file('/etc/environment', content, perms=0o644)
 
37
 
 
38
 
 
39
def all_relations(relation_types=CLIENT_RELATION_TYPES):
 
40
    for reltype in relation_types:
 
41
        for relid in hookenv.relation_ids(reltype):
 
42
            for unit in hookenv.related_units(relid):
 
43
                yield reltype, relid, unit, hookenv.relation_get(
 
44
                    unit=unit, rid=relid)
 
45
 
 
46
 
 
47
def rebuild_all_relations():
 
48
    config = hookenv.config()
 
49
 
 
50
    # Clear out old scripts and pgpass files
 
51
    if os.path.exists(SCRIPT_DIR):
 
52
        shutil.rmtree(SCRIPT_DIR)
 
53
    if os.path.exists(PGPASS_DIR):
 
54
        shutil.rmtree(PGPASS_DIR)
 
55
    host.mkdir(DATA_DIR, perms=0o755)
 
56
    host.mkdir(SCRIPT_DIR, perms=0o755)
 
57
    host.mkdir(PGPASS_DIR, group='ubuntu', perms=0o750)
 
58
 
 
59
    for _, relid, unit, relation in all_relations(relation_types=['db']):
 
60
        log("{} {} {!r}".format(relid, unit, relation), DEBUG)
 
61
 
 
62
        def_str = '<DEFAULT>'
 
63
        if config['database'] != relation.get('database', ''):
 
64
            log("Switching from database {} to {}".format(
 
65
                relation.get('database', '') or def_str,
 
66
                config['database'] or def_str), INFO)
 
67
 
 
68
        if config['roles'] != relation.get('roles', ''):
 
69
            log("Updating granted roles from {} to {}".format(
 
70
                relation.get('roles', '') or def_str,
 
71
                config['roles'] or def_str))
 
72
 
 
73
        hookenv.relation_set(
 
74
            relid, database=config['database'], roles=config['roles'])
 
75
 
 
76
        if 'user' in relation:
 
77
            rebuild_relation(relid, unit, relation)
 
78
 
 
79
    for _, relid, unit, relation in all_relations(relation_types=['db-admin']):
 
80
        log("{} {} {!r}".format(relid, unit, relation), DEBUG)
 
81
        if 'user' in relation:
 
82
            rebuild_relation(relid, unit, relation)
 
83
 
 
84
 
 
85
def rebuild_relation(relid, unit, relation):
 
86
    relname = relid.split(':')[0]
 
87
    unitname = unit.replace('/', '-')
 
88
    this_unit = hookenv.local_unit()
 
89
 
 
90
    allowed_units = relation.get('allowed-units', '')
 
91
    if this_unit not in allowed_units.split():
 
92
        log("Not yet authorized on {}".format(relid), INFO)
 
93
        return
 
94
 
 
95
    script_name = 'psql-{}-{}'.format(relname, unitname)
 
96
    build_script(script_name, relation)
 
97
    state = relation.get('state', None)
 
98
    if state in ('master', 'hot standby'):
 
99
        script_name = 'psql-{}-{}'.format(relname, state.replace(' ', '-'))
 
100
        build_script(script_name, relation)
 
101
 
 
102
 
 
103
def build_script(script_name, relation):
 
104
    # Install a wrapper to psql that connects it to the desired database
 
105
    # by default. One wrapper per unit per relation.
 
106
    script_path = os.path.abspath(os.path.join(SCRIPT_DIR, script_name))
 
107
    pgpass_path = os.path.abspath(os.path.join(PGPASS_DIR, script_name))
 
108
    script = dedent("""\
 
109
        #!/bin/sh
 
110
        exec env \\
 
111
            PGHOST={host} PGPORT={port} PGDATABASE={database} \\
 
112
            PGUSER={user} PGPASSFILE={pgpass} \\
 
113
            psql $@
 
114
        """).format(
 
115
            host=relation['host'],
 
116
            port=relation['port'],
 
117
            database=relation.get('database', ''),  # db-admin has no database
 
118
            user=relation['user'],
 
119
            pgpass=pgpass_path)
 
120
    log("Generating wrapper {}".format(script_path), INFO)
 
121
    host.write_file(
 
122
        script_path, script, owner="ubuntu", group="ubuntu", perms=0o700)
 
123
 
 
124
    # The wrapper requires access to the password, stored in a .pgpass
 
125
    # file so it isn't exposed in an environment variable or on the
 
126
    # command line.
 
127
    pgpass = "*:*:*:{user}:{password}".format(
 
128
        user=relation['user'], password=relation['password'])
 
129
    host.write_file(
 
130
        pgpass_path, pgpass, owner="ubuntu", group="ubuntu", perms=0o400)
 
131
 
 
132
 
 
133
hooks = hookenv.Hooks()
 
134
 
 
135
 
 
136
@hooks.hook()
 
137
def install():
 
138
    fetch.apt_install(
 
139
        ['language-pack-en', 'postgresql-client', 'python-psycopg2'],
 
140
        fatal=True)
 
141
    update_system_path()
 
142
 
 
143
 
 
144
@hooks.hook()
 
145
def upgrade_charm():
 
146
    # Per Bug #1205286, we can't store scripts and passwords in the
 
147
    # charm directory.
 
148
    if os.path.exists('bin'):
 
149
        shutil.rmtree('bin')
 
150
    if os.path.exists('pgpass'):
 
151
        shutil.rmtree('pgpass')
 
152
    update_system_path()
 
153
    return rebuild_all_relations()
 
154
 
 
155
 
 
156
@hooks.hook(
 
157
    'config-changed', 'db-admin-relation-broken',
 
158
    'db-admin-relation-changed', 'db-admin-relation-joined',
 
159
    'db-relation-broken', 'db-relation-changed', 'db-relation-joined')
 
160
def rebuild_hook():
 
161
    return rebuild_all_relations()
 
162
 
 
163
 
 
164
if __name__ == '__main__':
 
165
    hooks.execute(sys.argv)