hook.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. #! /usr/bin/env python
  2. # -*- encoding: utf-8 -*-
  3. # vim:fenc=utf-8:
  4. """Hooks are useful toy for fabric. Current versions of fabric do not
  5. support post action hooks."""
  6. import traceback
  7. from functools import wraps
  8. import mico
  9. import mico.output
  10. from __builtin__ import env
  11. # global variable for add_hooks()
  12. parent_task_name = ''
  13. def task_add_post_run_hook(hook, *args, **kwargs):
  14. '''Run hook after Fabric tasks have completed on all hosts
  15. Example usage:
  16. @add_post_run_hook(postrunfunc, 'arg1', 'arg2')
  17. def mytask():
  18. # ...
  19. '''
  20. def true_decorator(f):
  21. return add_hooks(post=hook, post_args=args, post_kwargs=kwargs)(f)
  22. return true_decorator
  23. def task_add_pre_run_hook(hook, *args, **kwargs):
  24. '''Run hook before Fabric tasks have completed on all hosts
  25. Example usage:
  26. @add_pre_run_hook(prerunfunc, 'arg1', 'arg2')
  27. def mytask():
  28. # ...
  29. '''
  30. def true_decorator(f):
  31. return add_hooks(pre=hook, pre_args=args, pre_kwargs=kwargs)(f)
  32. return true_decorator
  33. def add_hooks(pre=None, pre_args=(), pre_kwargs={},
  34. post=None, post_args=(), post_kwargs={}):
  35. '''
  36. Function decorator to be used with Fabric tasks. Adds pre-run
  37. and/or post-run hooks to a Fabric task. Uses env.all_hosts to
  38. determine when to run the post hook. Uses the global variable,
  39. parent_task_name, to check if the task is a subtask (i.e. a
  40. decorated task called by another decorated task). If it is a
  41. subtask, do not perform pre or post processing.
  42. pre: callable to be run before starting Fabric tasks
  43. pre_args: a tuple of arguments to be passed to "pre"
  44. pre_kwargs: a dict of keyword arguments to be passed to "pre"
  45. post: callable to be run after Fabric tasks have completed on all hosts
  46. post_args: a tuple of arguments to be passed to "post"
  47. post_kwargs: a dict of keyword arguments to be passed to "post"
  48. '''
  49. # create a namespace to save state across hosts and tasks
  50. class NS(object):
  51. run_counter = 0
  52. def true_decorator(f):
  53. @wraps(f)
  54. def f_wrapper(*args, **kwargs):
  55. # set state variables
  56. global parent_task_name
  57. if not parent_task_name:
  58. parent_task_name = f.__name__
  59. NS.run_counter += 1
  60. # pre-run processing
  61. if f.__name__ == parent_task_name and NS.run_counter == 1:
  62. if pre:
  63. pre(*pre_args, **pre_kwargs)
  64. # run the task
  65. r = None
  66. try:
  67. r = f(*args, **kwargs)
  68. except SystemExit:
  69. pass
  70. except mico.ExecutionError as e:
  71. mico.output.error(str(e))
  72. except:
  73. raise
  74. # post-run processing
  75. if (f.__name__ == parent_task_name and
  76. NS.run_counter >= len(env.all_hosts)):
  77. if post:
  78. post(*post_args, **post_kwargs)
  79. return r
  80. return f_wrapper
  81. return true_decorator
  82. from Queue import Queue
  83. def add_hook(context, action, parameters=()):
  84. """Enqueue one action with action parameters in environment"""
  85. if context in env:
  86. env[context].put((action,parameters))
  87. else:
  88. env[context] = Queue()
  89. env[context].put((action,parameters))
  90. def add_pre_hook(action, parameters=()):
  91. return add_hook("pre_hook", action, parameters)
  92. def add_post_hook(action, parameters=()):
  93. return add_hook("post_hook", action, parameters)
  94. def run_hook(context):
  95. if context in env:
  96. run = set()
  97. ret = []
  98. q = env[context]
  99. while not q.empty():
  100. action, parameters = q.get()
  101. if (action, parameters) in run:
  102. ## TODO: log skip running
  103. continue
  104. else:
  105. ret.append(action(*parameters))
  106. return ret
  107. else:
  108. return []
  109. def run_pre_hook():
  110. return run_hook("pre_hook")
  111. def run_post_hook():
  112. return run_hook("post_hook")