File: operation.py

package info (click to toggle)
pyinfra 0.2.2%2Bgit20161227.ec708ef-1
  • links: PTS, VCS
  • area: main
  • in suites: stretch
  • size: 11,804 kB
  • ctags: 677
  • sloc: python: 5,944; sh: 71; makefile: 11
file content (302 lines) | stat: -rw-r--r-- 10,768 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
# pyinfra
# File: pyinfra/api/operation.py
# Desc: wraps deploy script operations and puts commands -> pyinfra._ops

'''
Operations are the core of pyinfra. The ``@operation`` wrapper intercepts calls to the
function and instead diff against the remote server, outputting commands to the deploy
state. This is then run later by pyinfra's ``__main__`` or the :doc:`./api_operations`
module.
'''

from __future__ import unicode_literals

from os import path
from inspect import stack
from functools import wraps
from types import FunctionType

import six

from pyinfra import pseudo_state, pseudo_host
from pyinfra.pseudo_modules import PseudoModule

from .host import Host
from .state import State
from .attrs import wrap_attr_data
from .util import make_hash, get_arg_value, unroll_generators
from .exceptions import PyinfraError


class OperationMeta(object):
    def __init__(self, hash=None, commands=None):
        self.hash = hash
        self.commands = commands or []

        # Changed flag = did we do anything?
        self.changed = wrap_attr_data('changed', len(self.commands) > 0)


def add_op(state, op_func, *args, **kwargs):
    '''
    Prepare & add an operation to pyinfra.state by executing it on all hosts.

    Args:
        state (``pyinfra.api.State`` obj): the deploy state to add the operation to
        op_func (function): the operation function from one of the modules, ie \
            ``server.user``
        args/kwargs: passed to the operation function
    '''

    for host in state.inventory:
        op_func(state, host, *args, **kwargs)


def add_limited_op(state, op_func, hosts, *args, **kwargs):
    '''
    Prepare & add an operation to pyinfra.state by executing it on all hosts.

    Args:
        state (``pyinfra.api.State`` obj): the deploy state to add the operation to
        op_func (function): the operation function from one of the modules, ie \
            ``server.user``
        hosts (``pyinfra.api.Host`` obj or list): hosts this operation will be limited to
        args/kwargs: passed to the operation function
    '''

    if not isinstance(hosts, (list, tuple)):
        hosts = [hosts]

    # Set the limit
    state.limit_hosts = hosts

    # Add the op
    add_op(state, op_func, *args, **kwargs)

    # Remove the limit
    state.limit_hosts = []


def operation(func=None, pipeline_facts=None):
    '''
    Decorator that takes a simple module function and turn it into the internal operation
    representation that consists of a list of commands + options (sudo, (sudo|su)_user, env).
    '''

    # If not decorating, return a decorator which attaches any config to the function
    if func is None:
        def decorator(func):
            setattr(func, 'pipeline_facts', pipeline_facts)
            return operation(func)

        return decorator

    # Actually decorate!
    @wraps(func)
    def decorated_func(*args, **kwargs):
        # If we're in CLI mode, there's no state/host passed down, we need to use the
        # global "pseudo" modules.
        if len(args) < 2 or not (
            isinstance(args[0], (State, PseudoModule))
            and isinstance(args[1], (Host, PseudoModule))
        ):
            state = pseudo_state._module
            host = pseudo_host._module

            if not state.active:
                return OperationMeta()

            if state.in_op:
                raise PyinfraError(
                    'Nested operation called without state/host: {0}'.format(func)
                )

        # Otherwise (API mode) we just trim off the commands
        else:
            args_copy = list(args)
            state, host = args[0], args[1]
            args = args_copy[2:]

        # Pipelining? Now we have args, we can process the argspec and prep the pipe
        if state.pipelining:
            state.pipeline_facts.process(func, decorated_func, args, kwargs)

            # Not in op? Just drop the op into state.ops_to_pipeline and return here, this
            # will be re-run once the facts are gathered.
            if not state.in_op:
                state.ops_to_pipeline.append(
                    (host.name, decorated_func, args, kwargs.copy())
                )

        # Name the operation
        names = None
        autoname = False

        # Look for a set as the first argument
        if len(args) > 0 and isinstance(args[0], set):
            names = args[0]
            args_copy = list(args)
            args = args[1:]

        # Generate an operation name if needed (Module/Operation format)
        else:
            autoname = True
            module_bits = func.__module__.split('.')
            module_name = module_bits[-1]
            names = {'{0}/{1}'.format(module_name.title(), func.__name__.title())}

        # Locally & globally configurable
        sudo = kwargs.pop('sudo', state.config.SUDO)
        sudo_user = kwargs.pop('sudo_user', state.config.SUDO_USER)
        su_user = kwargs.pop('su_user', state.config.SU_USER)
        ignore_errors = kwargs.pop('ignore_errors', state.config.IGNORE_ERRORS)

        # Forces serial mode for this operation (--serial for one op)
        serial = kwargs.pop('serial', False)
        # Only runs this operation once
        run_once = kwargs.pop('run_once', False)
        # Timeout on running the command
        timeout = kwargs.pop('timeout', None)
        # Get a PTY before executing commands
        get_pty = kwargs.pop('get_pty', False)

        # Config env followed by command-level env
        env = state.config.ENV
        env.update(kwargs.pop('env', {}))

        # Get/generate a hash for this op
        op_hash = kwargs.pop('op', None)

        # If this op is being called inside another, just return here
        # (any unwanted/op-related kwargs removed above)
        if state.in_op:
            return func(state, host, *args, **kwargs) or []

        # Convert any AttrBase items (returned by host.data), see attrs.py.
        if op_hash is None:
            # Get the line number where this operation was called by looking through the
            # call stack for the first non-pyinfra line. This ensures two identical
            # operations (in terms of arguments/meta) will still generate two hashes.
            frames = stack()
            line_number = None

            for frame in frames:
                if not (
                    frame[3] in ('decorated_func', 'add_op', 'add_limited_op')
                    and frame[1].endswith(path.join('pyinfra', 'api', 'operation.py'))
                ):
                    line_number = frame[0].f_lineno
                    break

            op_hash = (
                names, sudo, sudo_user, su_user, line_number,
                ignore_errors, env, args, kwargs
            )

        op_hash = make_hash(op_hash)

        # Otherwise, flag as in-op and run it to get the commands
        state.in_op = True
        state.current_op_meta = (sudo, sudo_user, su_user, ignore_errors)

        # Generate actual arguments by parsing strings as jinja2 templates. This means
        # you can string format arguments w/o generating multiple operations. Only affects
        # top level operations, as must be run "in_op" so facts are gathered correctly.
        actual_args = [
            get_arg_value(state, host, arg)
            for arg in args
        ]

        actual_kwargs = {
            key: get_arg_value(state, host, arg)
            for key, arg in six.iteritems(kwargs)
        }

        # Convert to list as the result may be a generator
        commands = unroll_generators(func(state, host, *actual_args, **actual_kwargs))

        state.in_op = False
        state.current_op_meta = None

        # Make the operaton meta object for returning
        operation_meta = OperationMeta(op_hash, commands)

        # If we're pipelining, we don't actually want to add the operation as-is, just
        # collect the facts.
        if state.pipelining:
            return operation_meta

        # Add the hash to the operational order if not already in there. To ensure that
        # deploys run as defined in the deploy file *per host* we keep track of each hosts
        # latest op hash, and use that to insert new ones. Note that using the op kwarg
        # on operations can break the ordering when imbalanced (between if statements in
        # deploy file, for example).
        if op_hash not in state.op_order:
            previous_op_hash = state.meta[host.name]['latest_op_hash']

            if previous_op_hash:
                index = state.op_order.index(previous_op_hash)
            else:
                index = 0

            state.op_order.insert(index + 1, op_hash)

        state.meta[host.name]['latest_op_hash'] = op_hash

        # Run once and we've already added meta for this op? Stop here.
        if run_once and op_hash in state.op_meta:
            return operation_meta

        # Ensure shared (between servers) operation meta
        op_meta = state.op_meta.setdefault(op_hash, {
            'names': set(),
            'args': [],
            'su_user': su_user,
            'sudo': sudo,
            'sudo_user': sudo_user,
            'ignore_errors': ignore_errors,
            'serial': serial,
            'run_once': run_once,
            'timeout': timeout,
            'get_pty': get_pty,
        })

        # Add any new names to the set
        op_meta['names'].update(names)

        # Attach normal args, if we're auto-naming this operation
        if autoname:
            for arg in args:
                if isinstance(arg, FunctionType):
                    arg = arg.__name__

                if arg not in op_meta['args']:
                    op_meta['args'].append(arg)

            # Attach keyword args
            for key, value in six.iteritems(kwargs):
                arg = '='.join((str(key), str(value)))
                if arg not in op_meta['args']:
                    op_meta['args'].append(arg)

        # If we're limited, stop here - *after* we've created op_meta. This ensures the
        # meta object always exists, even if no hosts actually ever execute the op (due
        # to limit or otherwise).
        if state.limit_hosts is not None and host not in state.limit_hosts:
            return operation_meta

        # We're doing some commands, meta/ops++
        state.meta[host.name]['ops'] += 1
        state.meta[host.name]['commands'] += len(commands)

        # Add the server-relevant commands/env to the current server
        state.ops[host.name][op_hash] = {
            'commands': commands,
            'env': env,
        }

        # Return result meta for use in deploy scripts
        return operation_meta

    decorated_func._pyinfra_op = func
    return decorated_func