~adam-collard/charms/trusty/landscape-client/trunk

« back to all changes in this revision

Viewing changes to hooks/charmhelpers/core/hookenv.py

  • Committer: Christopher Glass
  • Date: 2014-06-12 11:41:32 UTC
  • mfrom: (37.2.3 landscape-client)
  • Revision ID: christopher.glass@canonical.com-20140612114132-dwvthvhhlg14wzjy
Merge changes from the ~landscape/charms/landscape-client/trunk branch.

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
"Interactions with the Juju environment"
 
2
# Copyright 2013 Canonical Ltd.
 
3
#
 
4
# Authors:
 
5
#  Charm Helpers Developers <juju@lists.ubuntu.com>
 
6
 
 
7
import os
 
8
import json
 
9
import yaml
 
10
import subprocess
 
11
import sys
 
12
import UserDict
 
13
from subprocess import CalledProcessError
 
14
 
 
15
CRITICAL = "CRITICAL"
 
16
ERROR = "ERROR"
 
17
WARNING = "WARNING"
 
18
INFO = "INFO"
 
19
DEBUG = "DEBUG"
 
20
MARKER = object()
 
21
 
 
22
cache = {}
 
23
 
 
24
 
 
25
def cached(func):
 
26
    """Cache return values for multiple executions of func + args
 
27
 
 
28
    For example:
 
29
 
 
30
        @cached
 
31
        def unit_get(attribute):
 
32
            pass
 
33
 
 
34
        unit_get('test')
 
35
 
 
36
    will cache the result of unit_get + 'test' for future calls.
 
37
    """
 
38
    def wrapper(*args, **kwargs):
 
39
        global cache
 
40
        key = str((func, args, kwargs))
 
41
        try:
 
42
            return cache[key]
 
43
        except KeyError:
 
44
            res = func(*args, **kwargs)
 
45
            cache[key] = res
 
46
            return res
 
47
    return wrapper
 
48
 
 
49
 
 
50
def flush(key):
 
51
    """Flushes any entries from function cache where the
 
52
    key is found in the function+args """
 
53
    flush_list = []
 
54
    for item in cache:
 
55
        if key in item:
 
56
            flush_list.append(item)
 
57
    for item in flush_list:
 
58
        del cache[item]
 
59
 
 
60
 
 
61
def log(message, level=None):
 
62
    """Write a message to the juju log"""
 
63
    command = ['juju-log']
 
64
    if level:
 
65
        command += ['-l', level]
 
66
    command += [message]
 
67
    subprocess.call(command)
 
68
 
 
69
 
 
70
class Serializable(UserDict.IterableUserDict):
 
71
    """Wrapper, an object that can be serialized to yaml or json"""
 
72
 
 
73
    def __init__(self, obj):
 
74
        # wrap the object
 
75
        UserDict.IterableUserDict.__init__(self)
 
76
        self.data = obj
 
77
 
 
78
    def __getattr__(self, attr):
 
79
        # See if this object has attribute.
 
80
        if attr in ("json", "yaml", "data"):
 
81
            return self.__dict__[attr]
 
82
        # Check for attribute in wrapped object.
 
83
        got = getattr(self.data, attr, MARKER)
 
84
        if got is not MARKER:
 
85
            return got
 
86
        # Proxy to the wrapped object via dict interface.
 
87
        try:
 
88
            return self.data[attr]
 
89
        except KeyError:
 
90
            raise AttributeError(attr)
 
91
 
 
92
    def __getstate__(self):
 
93
        # Pickle as a standard dictionary.
 
94
        return self.data
 
95
 
 
96
    def __setstate__(self, state):
 
97
        # Unpickle into our wrapper.
 
98
        self.data = state
 
99
 
 
100
    def json(self):
 
101
        """Serialize the object to json"""
 
102
        return json.dumps(self.data)
 
103
 
 
104
    def yaml(self):
 
105
        """Serialize the object to yaml"""
 
106
        return yaml.dump(self.data)
 
107
 
 
108
 
 
109
def execution_environment():
 
110
    """A convenient bundling of the current execution context"""
 
111
    context = {}
 
112
    context['conf'] = config()
 
113
    if relation_id():
 
114
        context['reltype'] = relation_type()
 
115
        context['relid'] = relation_id()
 
116
        context['rel'] = relation_get()
 
117
    context['unit'] = local_unit()
 
118
    context['rels'] = relations()
 
119
    context['env'] = os.environ
 
120
    return context
 
121
 
 
122
 
 
123
def in_relation_hook():
 
124
    """Determine whether we're running in a relation hook"""
 
125
    return 'JUJU_RELATION' in os.environ
 
126
 
 
127
 
 
128
def relation_type():
 
129
    """The scope for the current relation hook"""
 
130
    return os.environ.get('JUJU_RELATION', None)
 
131
 
 
132
 
 
133
def relation_id():
 
134
    """The relation ID for the current relation hook"""
 
135
    return os.environ.get('JUJU_RELATION_ID', None)
 
136
 
 
137
 
 
138
def local_unit():
 
139
    """Local unit ID"""
 
140
    return os.environ['JUJU_UNIT_NAME']
 
141
 
 
142
 
 
143
def remote_unit():
 
144
    """The remote unit for the current relation hook"""
 
145
    return os.environ['JUJU_REMOTE_UNIT']
 
146
 
 
147
 
 
148
def service_name():
 
149
    """The name service group this unit belongs to"""
 
150
    return local_unit().split('/')[0]
 
151
 
 
152
 
 
153
def hook_name():
 
154
    """The name of the currently executing hook"""
 
155
    return os.path.basename(sys.argv[0])
 
156
 
 
157
 
 
158
@cached
 
159
def config(scope=None):
 
160
    """Juju charm configuration"""
 
161
    config_cmd_line = ['config-get']
 
162
    if scope is not None:
 
163
        config_cmd_line.append(scope)
 
164
    config_cmd_line.append('--format=json')
 
165
    try:
 
166
        return json.loads(subprocess.check_output(config_cmd_line))
 
167
    except ValueError:
 
168
        return None
 
169
 
 
170
 
 
171
@cached
 
172
def relation_get(attribute=None, unit=None, rid=None):
 
173
    """Get relation information"""
 
174
    _args = ['relation-get', '--format=json']
 
175
    if rid:
 
176
        _args.append('-r')
 
177
        _args.append(rid)
 
178
    _args.append(attribute or '-')
 
179
    if unit:
 
180
        _args.append(unit)
 
181
    try:
 
182
        return json.loads(subprocess.check_output(_args))
 
183
    except ValueError:
 
184
        return None
 
185
    except CalledProcessError, e:
 
186
        if e.returncode == 2:
 
187
            return None
 
188
        raise
 
189
 
 
190
 
 
191
def relation_set(relation_id=None, relation_settings={}, **kwargs):
 
192
    """Set relation information for the current unit"""
 
193
    relation_cmd_line = ['relation-set']
 
194
    if relation_id is not None:
 
195
        relation_cmd_line.extend(('-r', relation_id))
 
196
    for k, v in (relation_settings.items() + kwargs.items()):
 
197
        if v is None:
 
198
            relation_cmd_line.append('{}='.format(k))
 
199
        else:
 
200
            relation_cmd_line.append('{}={}'.format(k, v))
 
201
    subprocess.check_call(relation_cmd_line)
 
202
    # Flush cache of any relation-gets for local unit
 
203
    flush(local_unit())
 
204
 
 
205
 
 
206
@cached
 
207
def relation_ids(reltype=None):
 
208
    """A list of relation_ids"""
 
209
    reltype = reltype or relation_type()
 
210
    relid_cmd_line = ['relation-ids', '--format=json']
 
211
    if reltype is not None:
 
212
        relid_cmd_line.append(reltype)
 
213
        return json.loads(subprocess.check_output(relid_cmd_line)) or []
 
214
    return []
 
215
 
 
216
 
 
217
@cached
 
218
def related_units(relid=None):
 
219
    """A list of related units"""
 
220
    relid = relid or relation_id()
 
221
    units_cmd_line = ['relation-list', '--format=json']
 
222
    if relid is not None:
 
223
        units_cmd_line.extend(('-r', relid))
 
224
    return json.loads(subprocess.check_output(units_cmd_line)) or []
 
225
 
 
226
 
 
227
@cached
 
228
def relation_for_unit(unit=None, rid=None):
 
229
    """Get the json represenation of a unit's relation"""
 
230
    unit = unit or remote_unit()
 
231
    relation = relation_get(unit=unit, rid=rid)
 
232
    for key in relation:
 
233
        if key.endswith('-list'):
 
234
            relation[key] = relation[key].split()
 
235
    relation['__unit__'] = unit
 
236
    return relation
 
237
 
 
238
 
 
239
@cached
 
240
def relations_for_id(relid=None):
 
241
    """Get relations of a specific relation ID"""
 
242
    relation_data = []
 
243
    relid = relid or relation_ids()
 
244
    for unit in related_units(relid):
 
245
        unit_data = relation_for_unit(unit, relid)
 
246
        unit_data['__relid__'] = relid
 
247
        relation_data.append(unit_data)
 
248
    return relation_data
 
249
 
 
250
 
 
251
@cached
 
252
def relations_of_type(reltype=None):
 
253
    """Get relations of a specific type"""
 
254
    relation_data = []
 
255
    reltype = reltype or relation_type()
 
256
    for relid in relation_ids(reltype):
 
257
        for relation in relations_for_id(relid):
 
258
            relation['__relid__'] = relid
 
259
            relation_data.append(relation)
 
260
    return relation_data
 
261
 
 
262
 
 
263
@cached
 
264
def relation_types():
 
265
    """Get a list of relation types supported by this charm"""
 
266
    charmdir = os.environ.get('CHARM_DIR', '')
 
267
    mdf = open(os.path.join(charmdir, 'metadata.yaml'))
 
268
    md = yaml.safe_load(mdf)
 
269
    rel_types = []
 
270
    for key in ('provides', 'requires', 'peers'):
 
271
        section = md.get(key)
 
272
        if section:
 
273
            rel_types.extend(section.keys())
 
274
    mdf.close()
 
275
    return rel_types
 
276
 
 
277
 
 
278
@cached
 
279
def relations():
 
280
    """Get a nested dictionary of relation data for all related units"""
 
281
    rels = {}
 
282
    for reltype in relation_types():
 
283
        relids = {}
 
284
        for relid in relation_ids(reltype):
 
285
            units = {local_unit(): relation_get(unit=local_unit(), rid=relid)}
 
286
            for unit in related_units(relid):
 
287
                reldata = relation_get(unit=unit, rid=relid)
 
288
                units[unit] = reldata
 
289
            relids[relid] = units
 
290
        rels[reltype] = relids
 
291
    return rels
 
292
 
 
293
 
 
294
@cached
 
295
def is_relation_made(relation, keys='private-address'):
 
296
    '''
 
297
    Determine whether a relation is established by checking for
 
298
    presence of key(s).  If a list of keys is provided, they
 
299
    must all be present for the relation to be identified as made
 
300
    '''
 
301
    if isinstance(keys, str):
 
302
        keys = [keys]
 
303
    for r_id in relation_ids(relation):
 
304
        for unit in related_units(r_id):
 
305
            context = {}
 
306
            for k in keys:
 
307
                context[k] = relation_get(k, rid=r_id,
 
308
                                          unit=unit)
 
309
            if None not in context.values():
 
310
                return True
 
311
    return False
 
312
 
 
313
 
 
314
def open_port(port, protocol="TCP"):
 
315
    """Open a service network port"""
 
316
    _args = ['open-port']
 
317
    _args.append('{}/{}'.format(port, protocol))
 
318
    subprocess.check_call(_args)
 
319
 
 
320
 
 
321
def close_port(port, protocol="TCP"):
 
322
    """Close a service network port"""
 
323
    _args = ['close-port']
 
324
    _args.append('{}/{}'.format(port, protocol))
 
325
    subprocess.check_call(_args)
 
326
 
 
327
 
 
328
@cached
 
329
def unit_get(attribute):
 
330
    """Get the unit ID for the remote unit"""
 
331
    _args = ['unit-get', '--format=json', attribute]
 
332
    try:
 
333
        return json.loads(subprocess.check_output(_args))
 
334
    except ValueError:
 
335
        return None
 
336
 
 
337
 
 
338
def unit_private_ip():
 
339
    """Get this unit's private IP address"""
 
340
    return unit_get('private-address')
 
341
 
 
342
 
 
343
class UnregisteredHookError(Exception):
 
344
    """Raised when an undefined hook is called"""
 
345
    pass
 
346
 
 
347
 
 
348
class Hooks(object):
 
349
    """A convenient handler for hook functions.
 
350
 
 
351
    Example:
 
352
        hooks = Hooks()
 
353
 
 
354
        # register a hook, taking its name from the function name
 
355
        @hooks.hook()
 
356
        def install():
 
357
            ...
 
358
 
 
359
        # register a hook, providing a custom hook name
 
360
        @hooks.hook("config-changed")
 
361
        def config_changed():
 
362
            ...
 
363
 
 
364
        if __name__ == "__main__":
 
365
            # execute a hook based on the name the program is called by
 
366
            hooks.execute(sys.argv)
 
367
    """
 
368
 
 
369
    def __init__(self):
 
370
        super(Hooks, self).__init__()
 
371
        self._hooks = {}
 
372
 
 
373
    def register(self, name, function):
 
374
        """Register a hook"""
 
375
        self._hooks[name] = function
 
376
 
 
377
    def execute(self, args):
 
378
        """Execute a registered hook based on args[0]"""
 
379
        hook_name = os.path.basename(args[0])
 
380
        if hook_name in self._hooks:
 
381
            self._hooks[hook_name]()
 
382
        else:
 
383
            raise UnregisteredHookError(hook_name)
 
384
 
 
385
    def hook(self, *hook_names):
 
386
        """Decorator, registering them as hooks"""
 
387
        def wrapper(decorated):
 
388
            for hook_name in hook_names:
 
389
                self.register(hook_name, decorated)
 
390
            else:
 
391
                self.register(decorated.__name__, decorated)
 
392
                if '_' in decorated.__name__:
 
393
                    self.register(
 
394
                        decorated.__name__.replace('_', '-'), decorated)
 
395
            return decorated
 
396
        return wrapper
 
397
 
 
398
 
 
399
def charm_dir():
 
400
    """Return the root directory of the current charm"""
 
401
    return os.environ.get('CHARM_DIR')