~niedbalski/charms/trusty/rabbitmq-server/fix-1378263

« back to all changes in this revision

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

  • Committer: James Page
  • Date: 2014-03-05 12:57:20 UTC
  • mto: This revision was merged to the branch mainline in revision 52.
  • Revision ID: james.page@canonical.com-20140305125720-fekpo1kpvh5tl93h
Redux take 1

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')