context_managers.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601
  1. """
  2. Context managers for use with the ``with`` statement.
  3. .. note:: When using Python 2.5, you will need to start your fabfile
  4. with ``from __future__ import with_statement`` in order to make use of
  5. the ``with`` statement (which is a regular, non ``__future__`` feature of
  6. Python 2.6+.)
  7. .. note:: If you are using multiple directly nested ``with`` statements, it can
  8. be convenient to use multiple context expressions in one single with
  9. statement. Instead of writing::
  10. with cd('/path/to/app'):
  11. with prefix('workon myvenv'):
  12. run('./manage.py syncdb')
  13. run('./manage.py loaddata myfixture')
  14. you can write::
  15. with cd('/path/to/app'), prefix('workon myvenv'):
  16. run('./manage.py syncdb')
  17. run('./manage.py loaddata myfixture')
  18. Note that you need Python 2.7+ for this to work. On Python 2.5 or 2.6, you
  19. can do the following::
  20. from contextlib import nested
  21. with nested(cd('/path/to/app'), prefix('workon myvenv')):
  22. ...
  23. Finally, note that `~fabric.context_managers.settings` implements
  24. ``nested`` itself -- see its API doc for details.
  25. """
  26. from contextlib import contextmanager, nested
  27. import socket
  28. import select
  29. from fabric.thread_handling import ThreadHandler
  30. from fabric.state import output, win32, connections, env
  31. from fabric import state
  32. from fabric.utils import isatty
  33. if not win32:
  34. import termios
  35. import tty
  36. def _set_output(groups, which):
  37. """
  38. Refactored subroutine used by ``hide`` and ``show``.
  39. """
  40. previous = {}
  41. try:
  42. # Preserve original values, pull in new given value to use
  43. for group in output.expand_aliases(groups):
  44. previous[group] = output[group]
  45. output[group] = which
  46. # Yield control
  47. yield
  48. finally:
  49. # Restore original values
  50. output.update(previous)
  51. def documented_contextmanager(func):
  52. wrapper = contextmanager(func)
  53. wrapper.undecorated = func
  54. return wrapper
  55. @documented_contextmanager
  56. def show(*groups):
  57. """
  58. Context manager for setting the given output ``groups`` to True.
  59. ``groups`` must be one or more strings naming the output groups defined in
  60. `~fabric.state.output`. The given groups will be set to True for the
  61. duration of the enclosed block, and restored to their previous value
  62. afterwards.
  63. For example, to turn on debug output (which is typically off by default)::
  64. def my_task():
  65. with show('debug'):
  66. run('ls /var/www')
  67. As almost all output groups are displayed by default, `show` is most useful
  68. for turning on the normally-hidden ``debug`` group, or when you know or
  69. suspect that code calling your own code is trying to hide output with
  70. `hide`.
  71. """
  72. return _set_output(groups, True)
  73. @documented_contextmanager
  74. def hide(*groups):
  75. """
  76. Context manager for setting the given output ``groups`` to False.
  77. ``groups`` must be one or more strings naming the output groups defined in
  78. `~fabric.state.output`. The given groups will be set to False for the
  79. duration of the enclosed block, and restored to their previous value
  80. afterwards.
  81. For example, to hide the "[hostname] run:" status lines, as well as
  82. preventing printout of stdout and stderr, one might use `hide` as follows::
  83. def my_task():
  84. with hide('running', 'stdout', 'stderr'):
  85. run('ls /var/www')
  86. """
  87. return _set_output(groups, False)
  88. @documented_contextmanager
  89. def _setenv(variables):
  90. """
  91. Context manager temporarily overriding ``env`` with given key/value pairs.
  92. A callable that returns a dict can also be passed. This is necessary when
  93. new values are being calculated from current values, in order to ensure that
  94. the "current" value is current at the time that the context is entered, not
  95. when the context manager is initialized. (See Issue #736.)
  96. This context manager is used internally by `settings` and is not intended
  97. to be used directly.
  98. """
  99. if callable(variables):
  100. variables = variables()
  101. clean_revert = variables.pop('clean_revert', False)
  102. previous = {}
  103. new = []
  104. for key, value in variables.iteritems():
  105. if key in state.env:
  106. previous[key] = state.env[key]
  107. else:
  108. new.append(key)
  109. state.env[key] = value
  110. try:
  111. yield
  112. finally:
  113. if clean_revert:
  114. for key, value in variables.iteritems():
  115. # If the current env value for this key still matches the
  116. # value we set it to beforehand, we are OK to revert it to the
  117. # pre-block value.
  118. if key in state.env and value == state.env[key]:
  119. if key in previous:
  120. state.env[key] = previous[key]
  121. else:
  122. del state.env[key]
  123. else:
  124. state.env.update(previous)
  125. for key in new:
  126. del state.env[key]
  127. def settings(*args, **kwargs):
  128. """
  129. Nest context managers and/or override ``env`` variables.
  130. `settings` serves two purposes:
  131. * Most usefully, it allows temporary overriding/updating of ``env`` with
  132. any provided keyword arguments, e.g. ``with settings(user='foo'):``.
  133. Original values, if any, will be restored once the ``with`` block closes.
  134. * The keyword argument ``clean_revert`` has special meaning for
  135. ``settings`` itself (see below) and will be stripped out before
  136. execution.
  137. * In addition, it will use `contextlib.nested`_ to nest any given
  138. non-keyword arguments, which should be other context managers, e.g.
  139. ``with settings(hide('stderr'), show('stdout')):``.
  140. .. _contextlib.nested: http://docs.python.org/library/contextlib.html#contextlib.nested
  141. These behaviors may be specified at the same time if desired. An example
  142. will hopefully illustrate why this is considered useful::
  143. def my_task():
  144. with settings(
  145. hide('warnings', 'running', 'stdout', 'stderr'),
  146. warn_only=True
  147. ):
  148. if run('ls /etc/lsb-release'):
  149. return 'Ubuntu'
  150. elif run('ls /etc/redhat-release'):
  151. return 'RedHat'
  152. The above task executes a `run` statement, but will warn instead of
  153. aborting if the ``ls`` fails, and all output -- including the warning
  154. itself -- is prevented from printing to the user. The end result, in this
  155. scenario, is a completely silent task that allows the caller to figure out
  156. what type of system the remote host is, without incurring the handful of
  157. output that would normally occur.
  158. Thus, `settings` may be used to set any combination of environment
  159. variables in tandem with hiding (or showing) specific levels of output, or
  160. in tandem with any other piece of Fabric functionality implemented as a
  161. context manager.
  162. If ``clean_revert`` is set to ``True``, ``settings`` will **not** revert
  163. keys which are altered within the nested block, instead only reverting keys
  164. whose values remain the same as those given. More examples will make this
  165. clear; below is how ``settings`` operates normally::
  166. # Before the block, env.parallel defaults to False, host_string to None
  167. with settings(parallel=True, host_string='myhost'):
  168. # env.parallel is True
  169. # env.host_string is 'myhost'
  170. env.host_string = 'otherhost'
  171. # env.host_string is now 'otherhost'
  172. # Outside the block:
  173. # * env.parallel is False again
  174. # * env.host_string is None again
  175. The internal modification of ``env.host_string`` is nullified -- not always
  176. desirable. That's where ``clean_revert`` comes in::
  177. # Before the block, env.parallel defaults to False, host_string to None
  178. with settings(parallel=True, host_string='myhost', clean_revert=True):
  179. # env.parallel is True
  180. # env.host_string is 'myhost'
  181. env.host_string = 'otherhost'
  182. # env.host_string is now 'otherhost'
  183. # Outside the block:
  184. # * env.parallel is False again
  185. # * env.host_string remains 'otherhost'
  186. Brand new keys which did not exist in ``env`` prior to using ``settings``
  187. are also preserved if ``clean_revert`` is active. When ``False``, such keys
  188. are removed when the block exits.
  189. .. versionadded:: 1.4.1
  190. The ``clean_revert`` kwarg.
  191. """
  192. managers = list(args)
  193. if kwargs:
  194. managers.append(_setenv(kwargs))
  195. return nested(*managers)
  196. def cd(path):
  197. """
  198. Context manager that keeps directory state when calling remote operations.
  199. Any calls to `run`, `sudo`, `get`, or `put` within the wrapped block will
  200. implicitly have a string similar to ``"cd <path> && "`` prefixed in order
  201. to give the sense that there is actually statefulness involved.
  202. .. note::
  203. `cd` only affects *remote* paths -- to modify *local* paths, use
  204. `~fabric.context_managers.lcd`.
  205. Because use of `cd` affects all such invocations, any code making use of
  206. those operations, such as much of the ``contrib`` section, will also be
  207. affected by use of `cd`.
  208. Like the actual 'cd' shell builtin, `cd` may be called with relative paths
  209. (keep in mind that your default starting directory is your remote user's
  210. ``$HOME``) and may be nested as well.
  211. Below is a "normal" attempt at using the shell 'cd', which doesn't work due
  212. to how shell-less SSH connections are implemented -- state is **not** kept
  213. between invocations of `run` or `sudo`::
  214. run('cd /var/www')
  215. run('ls')
  216. The above snippet will list the contents of the remote user's ``$HOME``
  217. instead of ``/var/www``. With `cd`, however, it will work as expected::
  218. with cd('/var/www'):
  219. run('ls') # Turns into "cd /var/www && ls"
  220. Finally, a demonstration (see inline comments) of nesting::
  221. with cd('/var/www'):
  222. run('ls') # cd /var/www && ls
  223. with cd('website1'):
  224. run('ls') # cd /var/www/website1 && ls
  225. .. note::
  226. This context manager is currently implemented by appending to (and, as
  227. always, restoring afterwards) the current value of an environment
  228. variable, ``env.cwd``. However, this implementation may change in the
  229. future, so we do not recommend manually altering ``env.cwd`` -- only
  230. the *behavior* of `cd` will have any guarantee of backwards
  231. compatibility.
  232. .. note::
  233. Space characters will be escaped automatically to make dealing with
  234. such directory names easier.
  235. .. versionchanged:: 1.0
  236. Applies to `get` and `put` in addition to the command-running
  237. operations.
  238. .. seealso:: `~fabric.context_managers.lcd`
  239. """
  240. return _change_cwd('cwd', path)
  241. def lcd(path):
  242. """
  243. Context manager for updating local current working directory.
  244. This context manager is identical to `~fabric.context_managers.cd`, except
  245. that it changes a different env var (`lcwd`, instead of `cwd`) and thus
  246. only affects the invocation of `~fabric.operations.local` and the local
  247. arguments to `~fabric.operations.get`/`~fabric.operations.put`.
  248. Relative path arguments are relative to the local user's current working
  249. directory, which will vary depending on where Fabric (or Fabric-using code)
  250. was invoked. You can check what this is with `os.getcwd
  251. <http://docs.python.org/release/2.6/library/os.html#os.getcwd>`_. It may be
  252. useful to pin things relative to the location of the fabfile in use, which
  253. may be found in :ref:`env.real_fabfile <real-fabfile>`
  254. .. versionadded:: 1.0
  255. """
  256. return _change_cwd('lcwd', path)
  257. def _change_cwd(which, path):
  258. path = path.replace(' ', '\ ')
  259. if state.env.get(which) and not path.startswith('/') and not path.startswith('~'):
  260. new_cwd = state.env.get(which) + '/' + path
  261. else:
  262. new_cwd = path
  263. return _setenv({which: new_cwd})
  264. def path(path, behavior='append'):
  265. """
  266. Append the given ``path`` to the PATH used to execute any wrapped commands.
  267. Any calls to `run` or `sudo` within the wrapped block will implicitly have
  268. a string similar to ``"PATH=$PATH:<path> "`` prepended before the given
  269. command.
  270. You may customize the behavior of `path` by specifying the optional
  271. ``behavior`` keyword argument, as follows:
  272. * ``'append'``: append given path to the current ``$PATH``, e.g.
  273. ``PATH=$PATH:<path>``. This is the default behavior.
  274. * ``'prepend'``: prepend given path to the current ``$PATH``, e.g.
  275. ``PATH=<path>:$PATH``.
  276. * ``'replace'``: ignore previous value of ``$PATH`` altogether, e.g.
  277. ``PATH=<path>``.
  278. .. note::
  279. This context manager is currently implemented by modifying (and, as
  280. always, restoring afterwards) the current value of environment
  281. variables, ``env.path`` and ``env.path_behavior``. However, this
  282. implementation may change in the future, so we do not recommend
  283. manually altering them directly.
  284. .. versionadded:: 1.0
  285. """
  286. return _setenv({'path': path, 'path_behavior': behavior})
  287. def prefix(command):
  288. """
  289. Prefix all wrapped `run`/`sudo` commands with given command plus ``&&``.
  290. This is nearly identical to `~fabric.operations.cd`, except that nested
  291. invocations append to a list of command strings instead of modifying a
  292. single string.
  293. Most of the time, you'll want to be using this alongside a shell script
  294. which alters shell state, such as ones which export or alter shell
  295. environment variables.
  296. For example, one of the most common uses of this tool is with the
  297. ``workon`` command from `virtualenvwrapper
  298. <http://www.doughellmann.com/projects/virtualenvwrapper/>`_::
  299. with prefix('workon myvenv'):
  300. run('./manage.py syncdb')
  301. In the above snippet, the actual shell command run would be this::
  302. $ workon myvenv && ./manage.py syncdb
  303. This context manager is compatible with `~fabric.context_managers.cd`, so
  304. if your virtualenv doesn't ``cd`` in its ``postactivate`` script, you could
  305. do the following::
  306. with cd('/path/to/app'):
  307. with prefix('workon myvenv'):
  308. run('./manage.py syncdb')
  309. run('./manage.py loaddata myfixture')
  310. Which would result in executions like so::
  311. $ cd /path/to/app && workon myvenv && ./manage.py syncdb
  312. $ cd /path/to/app && workon myvenv && ./manage.py loaddata myfixture
  313. Finally, as alluded to near the beginning,
  314. `~fabric.context_managers.prefix` may be nested if desired, e.g.::
  315. with prefix('workon myenv'):
  316. run('ls')
  317. with prefix('source /some/script'):
  318. run('touch a_file')
  319. The result::
  320. $ workon myenv && ls
  321. $ workon myenv && source /some/script && touch a_file
  322. Contrived, but hopefully illustrative.
  323. """
  324. return _setenv(lambda: {'command_prefixes': state.env.command_prefixes + [command]})
  325. @documented_contextmanager
  326. def char_buffered(pipe):
  327. """
  328. Force local terminal ``pipe`` be character, not line, buffered.
  329. Only applies on Unix-based systems; on Windows this is a no-op.
  330. """
  331. if win32 or not isatty(pipe):
  332. yield
  333. else:
  334. old_settings = termios.tcgetattr(pipe)
  335. tty.setcbreak(pipe)
  336. try:
  337. yield
  338. finally:
  339. termios.tcsetattr(pipe, termios.TCSADRAIN, old_settings)
  340. def shell_env(**kw):
  341. """
  342. Set shell environment variables for wrapped commands.
  343. For example, the below shows how you might set a ZeroMQ related environment
  344. variable when installing a Python ZMQ library::
  345. with shell_env(ZMQ_DIR='/home/user/local'):
  346. run('pip install pyzmq')
  347. As with `~fabric.context_managers.prefix`, this effectively turns the
  348. ``run`` command into::
  349. $ export ZMQ_DIR='/home/user/local' && pip install pyzmq
  350. Multiple key-value pairs may be given simultaneously.
  351. .. note::
  352. If used to affect the behavior of `~fabric.operations.local` when
  353. running from a Windows localhost, ``SET`` commands will be used to
  354. implement this feature.
  355. """
  356. return _setenv({'shell_env': kw})
  357. def _forwarder(chan, sock):
  358. # Bidirectionally forward data between a socket and a Paramiko channel.
  359. while True:
  360. r, w, x = select.select([sock, chan], [], [])
  361. if sock in r:
  362. data = sock.recv(1024)
  363. if len(data) == 0:
  364. break
  365. chan.send(data)
  366. if chan in r:
  367. data = chan.recv(1024)
  368. if len(data) == 0:
  369. break
  370. sock.send(data)
  371. chan.close()
  372. sock.close()
  373. @documented_contextmanager
  374. def remote_tunnel(remote_port, local_port=None, local_host="localhost",
  375. remote_bind_address="127.0.0.1"):
  376. """
  377. Create a tunnel forwarding a locally-visible port to the remote target.
  378. For example, you can let the remote host access a database that is
  379. installed on the client host::
  380. # Map localhost:6379 on the server to localhost:6379 on the client,
  381. # so that the remote 'redis-cli' program ends up speaking to the local
  382. # redis-server.
  383. with remote_tunnel(6379):
  384. run("redis-cli -i")
  385. The database might be installed on a client only reachable from the client
  386. host (as opposed to *on* the client itself)::
  387. # Map localhost:6379 on the server to redis.internal:6379 on the client
  388. with remote_tunnel(6379, local_host="redis.internal")
  389. run("redis-cli -i")
  390. ``remote_tunnel`` accepts up to four arguments:
  391. * ``remote_port`` (mandatory) is the remote port to listen to.
  392. * ``local_port`` (optional) is the local port to connect to; the default is
  393. the same port as the remote one.
  394. * ``local_host`` (optional) is the locally-reachable computer (DNS name or
  395. IP address) to connect to; the default is ``localhost`` (that is, the
  396. same computer Fabric is running on).
  397. * ``remote_bind_address`` (optional) is the remote IP address to bind to
  398. for listening, on the current target. It should be an IP address assigned
  399. to an interface on the target (or a DNS name that resolves to such IP).
  400. You can use "0.0.0.0" to bind to all interfaces.
  401. .. note::
  402. By default, most SSH servers only allow remote tunnels to listen to the
  403. localhost interface (127.0.0.1). In these cases, `remote_bind_address`
  404. is ignored by the server, and the tunnel will listen only to 127.0.0.1.
  405. .. versionadded: 1.6
  406. """
  407. if local_port is None:
  408. local_port = remote_port
  409. sockets = []
  410. channels = []
  411. threads = []
  412. def accept(channel, (src_addr, src_port), (dest_addr, dest_port)):
  413. channels.append(channel)
  414. sock = socket.socket()
  415. sockets.append(sock)
  416. try:
  417. sock.connect((local_host, local_port))
  418. except Exception:
  419. print "[%s] rtunnel: cannot connect to %s:%d (from local)" % (env.host_string, local_host, local_port)
  420. channel.close()
  421. return
  422. print "[%s] rtunnel: opened reverse tunnel: %r -> %r -> %r"\
  423. % (env.host_string, channel.origin_addr,
  424. channel.getpeername(), (local_host, local_port))
  425. th = ThreadHandler('fwd', _forwarder, channel, sock)
  426. threads.append(th)
  427. transport = connections[env.host_string].get_transport()
  428. transport.request_port_forward(remote_bind_address, remote_port, handler=accept)
  429. try:
  430. yield
  431. finally:
  432. for sock, chan, th in zip(sockets, channels, threads):
  433. sock.close()
  434. chan.close()
  435. th.thread.join()
  436. th.raise_if_needed()
  437. transport.cancel_port_forward(remote_bind_address, remote_port)
  438. quiet = lambda: settings(hide('everything'), warn_only=True)
  439. quiet.__doc__ = """
  440. Alias to ``settings(hide('everything'), warn_only=True)``.
  441. Useful for wrapping remote interrogative commands which you expect to fail
  442. occasionally, and/or which you want to silence.
  443. Example::
  444. with quiet():
  445. have_build_dir = run("test -e /tmp/build").succeeded
  446. When used in a task, the above snippet will not produce any ``run: test -e
  447. /tmp/build`` line, nor will any stdout/stderr display, and command failure
  448. is ignored.
  449. .. seealso::
  450. :ref:`env.warn_only <warn_only>`,
  451. `~fabric.context_managers.settings`,
  452. `~fabric.context_managers.hide`
  453. .. versionadded:: 1.5
  454. """
  455. warn_only = lambda: settings(warn_only=True)
  456. warn_only.__doc__ = """
  457. Alias to ``settings(warn_only=True)``.
  458. .. seealso::
  459. :ref:`env.warn_only <warn_only>`,
  460. `~fabric.context_managers.settings`,
  461. `~fabric.context_managers.quiet`
  462. """