-
Notifications
You must be signed in to change notification settings - Fork 303
/
wrap.py
396 lines (328 loc) · 14.4 KB
/
wrap.py
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
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
'''
Wraps shell commands and sends the result to Datadog as events. Ex:
dogwrap -n test-job -k $API_KEY --submit_mode all "ls -lah"
Note that you need to enclose your command in quotes to prevent python
from thinking the command line arguments belong to the python command
instead of the wrapped command.
You can also have the script only send events if they fail:
dogwrap -n test-job -k $API_KEY --submit_mode errors "ls -lah"
And you can give the command a timeout too:
dogwrap -n test-job -k $API_KEY --timeout=1 "sleep 3"
'''
# stdlib
from __future__ import print_function
from copy import copy
import optparse
import subprocess
import sys
import threading
import time
# datadog
from datadog import initialize, api
from datadog.util.config import get_version
from datadog.util.compat import is_p3k
SUCCESS = 'success'
ERROR = 'error'
WARNING = 'warning'
MAX_EVENT_BODY_LENGTH = 3000
class Timeout(Exception):
pass
class OutputReader(threading.Thread):
'''
Thread collecting the output of a subprocess, optionally forwarding it to
a given file descriptor and storing it for further retrieval.
'''
def __init__(self, proc_out, fwd_out=None):
'''
Instantiates an OutputReader.
:param proc_out: the output to read
:type proc_out: file descriptor
:param fwd_out: the output to forward to (None to disable forwarding)
:type fwd_out: file descriptor or None
'''
threading.Thread.__init__(self)
self.daemon = True
self._out_content = b""
self._out = proc_out
self._fwd_out = fwd_out
def run(self):
'''
Thread's main loop: collects the output optionnally forwarding it to
the file descriptor passed in the constructor.
'''
for line in iter(self._out.readline, b''):
if self._fwd_out is not None:
self._fwd_out.write(line)
self._out_content += line
self._out.close()
@property
def content(self):
'''
The content stored in out so far. (Not threadsafe, wait with .join())
'''
return self._out_content
def poll_proc(proc, sleep_interval, timeout):
'''
Polls the process until it returns or a given timeout has been reached
'''
start_time = time.time()
returncode = None
while returncode is None:
returncode = proc.poll()
if time.time() - start_time > timeout:
raise Timeout()
else:
time.sleep(sleep_interval)
return returncode
def execute(cmd, cmd_timeout, sigterm_timeout, sigkill_timeout,
proc_poll_interval, buffer_outs):
'''
Launches the process and monitors its outputs
'''
start_time = time.time()
returncode = -1
stdout = b''
stderr = b''
try:
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
except Exception:
print(u"Failed to execute %s" % (repr(cmd)), file=sys.stderr)
raise
try:
# Let's that the threads collecting the output from the command in the
# background
stdout = sys.stdout.buffer if is_p3k() else sys.stdout
stderr = sys.stderr.buffer if is_p3k() else sys.stderr
out_reader = OutputReader(proc.stdout, stdout if not buffer_outs else None)
err_reader = OutputReader(proc.stderr, stderr if not buffer_outs else None)
out_reader.start()
err_reader.start()
# Let's quietly wait from the program's completion here et get the exit
# code when it finishes
returncode = poll_proc(proc, proc_poll_interval, cmd_timeout)
# Let's harvest the outputs collected by our background threads
# after making sure they're done reading it.
out_reader.join()
err_reader.join()
stdout = out_reader.content
stderr = err_reader.content
duration = time.time() - start_time
except Timeout:
duration = time.time() - start_time
try:
proc.terminate()
sigterm_start = time.time()
try:
print("Command timed out after %.2fs, killing with SIGTERM", file=sys.stderr) \
% (time.time() - start_time)
poll_proc(proc, proc_poll_interval, sigterm_timeout)
returncode = Timeout
except Timeout:
print("SIGTERM timeout failed after %.2fs, killing with SIGKILL", file=sys.stderr) \
% (time.time() - sigterm_start)
proc.kill()
poll_proc(proc, proc_poll_interval, sigkill_timeout)
returncode = Timeout
except OSError as e:
# Ignore OSError 3: no process found.
if e.errno != 3:
raise
return returncode, stdout, stderr, duration
def trim_text(text, max_len):
"""
Trim input text to fit the `max_len` condition.
If trim is needed: keep the first 1/3rd of the budget on the top,
and the other 2 thirds on the bottom.
"""
if len(text) <= max_len:
return text
trimmed_text = \
u"{top_third}\n"\
u"```\n" \
u"*...trimmed...*\n" \
u"```\n" \
u"{bottom_two_third}\n".format(
top_third=text[:max_len // 3],
bottom_two_third=text[len(text) - (2 * max_len) // 3:]
)
return trimmed_text
def build_event_body(cmd, returncode, stdout, stderr, notifications):
"""
Format and return an event body.
Note: do not exceed MAX_EVENT_BODY_LENGTH length.
"""
fmt_stdout = u""
fmt_stderr = u""
fmt_notifications = u""
max_length = MAX_EVENT_BODY_LENGTH // 2 if stdout and stderr else MAX_EVENT_BODY_LENGTH
if stdout:
fmt_stdout = u"**>>>> STDOUT <<<<**\n```\n{stdout} \n```\n".format(
stdout=trim_text(stdout.decode("utf-8", "replace"), max_length)
)
if stderr:
fmt_stderr = u"**>>>> STDERR <<<<**\n```\n{stderr} \n```\n".format(
stderr=trim_text(stderr.decode("utf-8", "replace"), max_length)
)
if notifications:
fmt_notifications = u"**>>>> NOTIFICATIONS <<<<**\n\n {notifications}\n".format(
notifications=notifications.decode("utf-8", "replace")
)
return \
u"%%%\n" \
u"**>>>> CMD <<<<**\n```\n{command} \n```\n" \
u"**>>>> EXIT CODE <<<<**\n\n {returncode}\n\n\n" \
u"{stdout}" \
u"{stderr}" \
u"{notifications}" \
u"%%%\n".format(
command=cmd,
returncode=returncode,
stdout=fmt_stdout,
stderr=fmt_stderr,
notifications=fmt_notifications,
)
def generate_warning_codes(option, opt, options_warning):
try:
# options_warning is a string e.g.: --warning_codes 123,456,789
# we need to create a list from it
warning_codes = options_warning.split(",")
return warning_codes
except ValueError:
raise optparse.OptionValueError("option %s: invalid warning codes value(s): %r" % (opt, options_warning))
class DogwrapOption(optparse.Option):
# https://docs.python.org/3.7/library/optparse.html#adding-new-types
TYPES = optparse.Option.TYPES + ("warning_codes",)
TYPE_CHECKER = copy(optparse.Option.TYPE_CHECKER)
TYPE_CHECKER["warning_codes"] = generate_warning_codes
def parse_options(raw_args=None):
'''
Parse the raw command line options into an options object and the remaining command string
'''
parser = optparse.OptionParser(usage="%prog -n [event_name] -k [api_key] --submit_mode \
[ all | errors | warnings] [options] \"command\". \n\nNote that you need to enclose your command in \
quotes to prevent python executing as soon as there is a space in your command. \n \nNOTICE: In \
normal mode, the whole stderr is printed before stdout, in flush_live mode they will be mixed but \
there is not guarantee that messages sent by the command on both stderr and stdout are printed in \
the order they were sent.", version="%prog {0}".format(get_version()), option_class=DogwrapOption)
parser.add_option('-n', '--name', action='store', type='string', help="the name of the event \
as it should appear on your Datadog stream")
parser.add_option('-k', '--api_key', action='store', type='string',
help="your DataDog API Key")
parser.add_option('-s', '--site', action='store', type='choice', default='us', choices=['us', 'eu'], help="The site \
to send data, US or EU, default: US")
parser.add_option('-m', '--submit_mode', action='store', type='choice',
default='errors', choices=['errors', 'warnings', 'all'], help="[ all | errors | warnings ] if set \
to error, an event will be sent only of the command exits with a non zero exit status or if it \
times out. If set to warning, a list of exit codes need to be provided")
parser.add_option('--warning_codes', action='store', type='warning_codes', dest='warning_codes',
help="comma separated list of warning codes, e.g: 127,255")
parser.add_option('-p', '--priority', action='store', type='choice', choices=['normal', 'low'],
help="the priority of the event (default: 'normal')")
parser.add_option('-t', '--timeout', action='store', type='int', default=60 * 60 * 24,
help="(in seconds) a timeout after which your command must be aborted. An \
event will be sent to your DataDog stream (default: 24hours)")
parser.add_option('--sigterm_timeout', action='store', type='int', default=60 * 2,
help="(in seconds) When your command times out, the \
process it triggers is sent a SIGTERM. If this sigterm_timeout is reached, it will be sent a \
SIGKILL signal. (default: 2m)")
parser.add_option('--sigkill_timeout', action='store', type='int', default=60,
help="(in seconds) how long to wait at most after SIGKILL \
has been sent (default: 60s)")
parser.add_option('--proc_poll_interval', action='store', type='float', default=0.5,
help="(in seconds). interval at which your command will be polled \
(default: 500ms)")
parser.add_option('--notify_success', action='store', type='string', default='',
help="a message string and @people directives to send notifications in \
case of success.")
parser.add_option('--notify_error', action='store', type='string', default='',
help="a message string and @people directives to send notifications in \
case of error.")
parser.add_option('--notify_warning', action='store', type='string', default='',
help="a message string and @people directives to send notifications in \
case of warning.")
parser.add_option('-b', '--buffer_outs', action='store_true', dest='buffer_outs', default=False,
help="displays the stderr and stdout of the command only once it has \
returned (the command outputs remains buffered in dogwrap meanwhile)")
parser.add_option('--tags', action='store', type='string', dest='tags', default='',
help="comma separated list of tags")
options, args = parser.parse_args(args=raw_args)
if is_p3k():
cmd = ' '.join(args)
else:
cmd = b' '.join(args).decode('utf-8')
return options, cmd
def main():
options, cmd = parse_options()
# If silent is checked we force the outputs to be buffered (and therefore
# not forwarded to the Terminal streams) and we just avoid printing the
# buffers at the end
returncode, stdout, stderr, duration = execute(
cmd, options.timeout,
options.sigterm_timeout, options.sigkill_timeout,
options.proc_poll_interval, options.buffer_outs)
if options.site == 'eu':
api_host = 'https://api.datadoghq.eu'
else:
api_host = 'https://api.datadoghq.com'
initialize(api_key=options.api_key, api_host=api_host)
host = api._host_name
warning_codes = None
if options.warning_codes:
# Convert warning codes from string to int since return codes will evaluate the latter
warning_codes = list(map(int, options.warning_codes))
if returncode == 0:
alert_type = SUCCESS
event_priority = 'low'
event_title = u'[%s] %s succeeded in %.2fs' % (host, options.name,
duration)
elif returncode != 0 and options.submit_mode == 'warnings':
if not warning_codes:
# the list of warning codes is empty - the option was not specified
print("A comma separated list of exit codes need to be provided")
sys.exit()
elif returncode in warning_codes:
alert_type = WARNING
event_priority = 'normal'
event_title = u'[%s] %s failed in %.2fs' % (host, options.name,
duration)
else:
print("Command exited with a different exit code that the one(s) provided")
sys.exit()
else:
alert_type = ERROR
event_priority = 'normal'
if returncode is Timeout:
event_title = u'[%s] %s timed out after %.2fs' % (host, options.name, duration)
returncode = -1
else:
event_title = u'[%s] %s failed in %.2fs' % (host, options.name, duration)
notifications = ""
if alert_type == SUCCESS and options.notify_success:
notifications = options.notify_success
elif alert_type == ERROR and options.notify_error:
notifications = options.notify_error
elif alert_type == WARNING and options.notify_warning:
notifications = options.notify_warning
if options.tags:
tags = [t.strip() for t in options.tags.split(',')]
else:
tags = None
event_body = build_event_body(cmd, returncode, stdout, stderr, notifications)
event = {
'alert_type': alert_type,
'aggregation_key': options.name,
'host': host,
'priority': options.priority or event_priority,
'tags': tags
}
if options.buffer_outs:
if is_p3k():
stderr = stderr.decode('utf-8')
stdout = stdout.decode('utf-8')
print(stderr.strip(), file=sys.stderr)
print(stdout.strip(), file=sys.stdout)
if options.submit_mode == 'all' or returncode != 0:
api.Event.create(title=event_title, text=event_body, **event)
sys.exit(returncode)
if __name__ == '__main__':
main()