def setup_limits(conf_file,
limits_file,
do_reload=True,
dry_run=False,
debug=False):
"""
Set up or update limits in the Redis database.
:param conf_file: Name of the configuration file, for connecting
to the Redis database.
:param limits_file: Name of the XML file describing the limits to
configure.
:param do_reload: Controls reloading behavior. If True (the
default), a reload command is issued. If False,
no reload command is issued. String values
result in a reload command of the given load
type, and integer or float values result in a
reload command of type 'spread' with the given
spread interval.
:param dry_run: If True, no changes are made to the database.
Implies debug=True.
:param debug: If True, debugging messages are emitted while
loading the limits and updating the database.
"""
# If dry_run is set, default debug to True
if dry_run:
debug = True
# Connect to the database...
conf = config.Config(conf_file=conf_file)
db = conf.get_database()
limits_key = conf['control'].get('limits_key', 'limits')
control_channel = conf['control'].get('channel', 'control')
# Parse the limits file
limits_tree = etree.parse(limits_file)
# Now, we parse the limits XML file
lims = []
for idx, lim in enumerate(limits_tree.getroot()):
# Skip tags we don't recognize
if lim.tag != 'limit':
warnings.warn("Unrecognized tag %r in limits file at index %d" %
(lim.tag, idx))
continue
# Construct the limit and add it to the list of limits
try:
lims.append(parse_limit_node(db, idx, lim))
except Exception as exc:
warnings.warn("Couldn't understand limit at index %d: %s" %
(idx, exc))
continue
# Now that we have the limits, let's install them
if debug:
print >> sys.stderr, "Installing the following limits:"
for lim in lims:
print >> sys.stderr, " %r" % lim
if not dry_run:
database.limit_update(db, limits_key, lims)
# Were we requested to reload the limits?
if do_reload is False:
return
# OK, figure out what kind of reload to do
params = []
if do_reload is True:
# Nothing to do; use default semantics
pass
elif (isinstance(do_reload, (int, long, float))
or (isinstance(do_reload, basestring) and do_reload.isdigit())):
params = ['spread', do_reload]
else:
params = [str(do_reload)]
# Issue the reload command
if debug:
cmd = ['reload']
cmd.extend(params)
print >> sys.stderr, ("Issuing command: %s" %
' '.join(str(c) for c in cmd))
if not dry_run:
database.command(db, control_channel, 'reload', *params)
def turnstile_command(conf_file,
command,
arguments=[],
channel=None,
debug=False):
"""
Issue a command to all running control daemons.
:param conf_file: Name of the configuration file.
:param command: The command to execute. Note that 'ping' is
handled specially; in particular, the "channel"
parameter is implied. (A random value will be
used for the channel to listen on.)
:param arguments: A list of arguments for the command. Note that
the colon character (':') cannot be used.
:param channel: If not None, specifies the name of a message
channel to listen for responses on. Will wait
indefinitely; to terminate the listening loop, use
the keyboard interrupt sequence.
:param debug: If True, debugging messages are emitted while
sending the command.
"""
# Connect to the database...
conf = config.Config(conf_file=conf_file)
db = conf.get_database()
control_channel = conf['control'].get('channel', 'control')
# Now, set up the command
command = command.lower()
ts_conv = False
if command == 'ping':
# We handle 'ping' specially; first, figure out the channel
if arguments:
channel = arguments[0]
else:
channel = str(uuid.uuid4())
arguments = [channel]
# Next, add on a timestamp
if len(arguments) < 2:
arguments.append(time.time())
ts_conv = True
# Limit the argument list length
arguments = arguments[:2]
# OK, the command is all set up. Let us now send the command...
if debug:
cmd = [command] + arguments
print >> sys.stderr, ("Issuing command: %s" % ' '.join(cmd))
database.command(db, control_channel, command, *arguments)
# Were we asked to listen on a channel?
if not channel:
return
# OK, let's subscribe to the channel...
pubsub = db.pubsub()
pubsub.subscribe(channel)
# Now we listen...
try:
count = 0
for msg in pubsub.listen():
# Make sure the message is one we're interested in
if debug:
formatted = pprint.pformat(msg)
print >> sys.stderr, "Received message: %s" % formatted
if (msg['type'] not in ('pmessage', 'message')
or msg['channel'] != channel):
continue
count += 1
# Figure out the response
response = msg['data'].split(':')
# If this is a 'pong' and ts_conv is true, add an RTT to
# the response
if ts_conv and response[0] == 'pong':
try:
rtt = (time.time() - float(response[2])) * 100
response.append('(RTT %.2fms)' % rtt)
except Exception:
# IndexError or ValueError, probably; ignore it
pass
# Print out the response
print "Response % 5d: %s" % (count, ' '.join(response))
except KeyboardInterrupt:
# We want to break out of the loop, but not return any error
# to the caller...
pass
def test_command(self):
db = mock.Mock()
database.command(db, 'channel', 'command', 'one', 2, 3.14)
db.publish.assert_called_once_with('channel', 'command:one:2:3.14')