scripts/write_live_config

   1 #!/usr/bin/env python
   2 # The contents of this file are subject to the Common Public Attribution
   3 # License Version 1.0. (the "License"); you may not use this file except in
   4 # compliance with the License. You may obtain a copy of the License at
   5 # http://code.reddit.com/LICENSE. The License is based on the Mozilla Public
   6 # License Version 1.1, but Sections 14 and 15 have been added to cover use of
   7 # software over a computer network and provide for limited attribution for the
   8 # Original Developer. In addition, Exhibit A has been modified to be consistent
   9 # with Exhibit B.
  10 #
  11 # Software distributed under the License is distributed on an "AS IS" basis,
  12 # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
  13 # the specific language governing rights and limitations under the License.
  14 #
  15 # The Original Code is reddit.
  16 #
  17 # The Original Developer is the Initial Developer.  The Initial Developer of
  18 # the Original Code is reddit Inc.
  19 #
  20 # All portions of the code written by reddit are Copyright (c) 2006-2015 reddit
  21 # Inc. All Rights Reserved.
  22 ###############################################################################
  23 """Read config from an INI file and put it in ZooKeeper for instant use."""
  24
  25 from difflib import SequenceMatcher
  26 import itertools
  27 import os
  28 from pprint import PrettyPrinter
  29 import sys
  30 import json
  31 import getpass
  32 import ConfigParser
  33 import zlib
  34
  35 import kazoo.client
  36 from kazoo.security import make_acl, make_digest_acl
  37 from kazoo.exceptions import NoAuthException
  38
  39 from r2.lib.zookeeper import LiveConfig, connect_to_zookeeper
  40 from r2.lib.app_globals import extract_live_config, LIVE_CONFIG_NODE
  41 from r2.lib.configparse import ConfigValue
  42 from r2.lib.plugin import PluginLoader
  43 from r2.lib.utils import parse_ini_file
  44
  45
  46 USERNAME = "live-config"
  47
  48
  49 def write_config_to_zookeeper(node, username, password, config, live_config):
  50     """Write given configuration to ZooKeeper with correct security etc."""
  51
  52     # read the zk configuration from the app's config
  53     zk_hostlist = config.get("DEFAULT", "zookeeper_connection_string")
  54     app_username = config.get("DEFAULT", "zookeeper_username")
  55     app_password = config.get("DEFAULT", "zookeeper_password")
  56
  57     # connect to zk!
  58     client = connect_to_zookeeper(zk_hostlist, (username, password))
  59
  60     # ensure that the path leading up to the config node exists. if it doesn't,
  61     # create it with ACLs such that new stuff can be added below it, but no one
  62     # but we can delete nodes.
  63     parent_path = os.path.dirname(node)
  64     client.ensure_path(parent_path, acl=[
  65         # only we can delete children
  66         make_digest_acl(username, password, delete=True),
  67
  68         # anyone authenticated can read/list children/create children
  69         make_acl("auth", "", read=True, create=True),
  70     ])
  71
  72     # create or update the config node ensuring that only we can write to it.
  73     json_data = json.dumps(live_config)
  74     compressed_data = "gzip" + zlib.compress(json_data)
  75
  76     try:
  77         client.create(node, compressed_data, acl=[
  78             make_digest_acl(username, password, read=True, write=True),
  79             make_digest_acl(app_username, app_password, read=True),
  80         ])
  81     except kazoo.exceptions.NodeExistsException:
  82         client.set(node, compressed_data)
  83
  84
  85 def get_comparable_repr(obj):
  86     """Return a representation of the object that can be string-compared."""
  87
  88     # If the object is a dict, we'll use the pprint module, because it
  89     # automatically sorts by key when generating its output for dicts
  90     if isinstance(obj, dict):
  91         # specify a huge width so it never tries to wrap the output
  92         printer = PrettyPrinter(width=1000000)
  93         return printer.pformat(obj)
  94
  95     # Otherwise, just use the standard repr for that object type
  96     return repr(obj)
  97
  98
  99 def print_dict_diff(old, new):
 100     """Output changes between two dicts."""
 101
 102     old_keys = set(old.keys())
 103     new_keys = set(new.keys())
 104
 105     # figure out which keys are interesting that we need to output
 106     removed_keys = old_keys - new_keys
 107     added_keys = new_keys - old_keys
 108     both_keys = old_keys & new_keys
 109     changed_keys = [key for key in both_keys if new[key] != old[key]]
 110
 111     for key in changed_keys[:]:
 112         # if the value type changed, or it's a type with short values,
 113         # just display it as a removal and addition
 114         if (type(new[key]) != type(old[key]) or
 115                 isinstance(new[key], (int, float, bool))):
 116             added_keys.add(key)
 117             removed_keys.add(key)
 118             changed_keys.remove(key)
 119             continue
 120
 121         # otherwise, see how similar the reprs are, and if it's less
 122         # than 50% similar, just display it as a removal and addition
 123         old_repr = get_comparable_repr(old[key])
 124         new_repr = get_comparable_repr(new[key])
 125         matcher = SequenceMatcher(a=old_repr, b=new_repr)
 126
 127         if matcher.ratio() < 0.5:
 128             added_keys.add(key)
 129             removed_keys.add(key)
 130             changed_keys.remove(key)
 131             continue
 132
 133     keys = sorted(set(itertools.chain(removed_keys, added_keys, changed_keys)))
 134     if not keys:
 135         print "No changes found. Did you forget to deploy the INI change?"
 136         return
 137
 138     max_key_length = max(len(key) for key in keys)
 139
 140     for key in keys:
 141         if key in removed_keys:
 142             print "- {key:<{length}s} = {value!r}".format(
 143                 key=key, value=old[key], length=max_key_length)
 144
 145         if key in added_keys:
 146             print "+ {key:<{length}s} = {value!r}".format(
 147                 key=key, value=new[key], length=max_key_length)
 148
 149         if key in changed_keys:
 150             print "! {key:<{length}s} :".format(key=key, length=max_key_length)
 151
 152             old_repr = get_comparable_repr(old[key])
 153             new_repr = get_comparable_repr(new[key])
 154             matcher = SequenceMatcher(a=old_repr, b=new_repr)
 155
 156             for tag, i, j, m, n in matcher.get_opcodes():
 157                 if tag == "equal":
 158                     continue
 159                 elif tag == "replace":
 160                     print '    ! "{}" to "{}"'.format(
 161                         matcher.a[i:j], matcher.b[m:n])
 162                 elif tag == "insert":
 163                     print '    + "{}"'.format(matcher.b[m:n])
 164                 elif tag == "delete":
 165                     print '    - "{}"'.format(matcher.a[i:j])
 166
 167
 168 def get_current_live_config(config):
 169     """Return the current live config values as a dict."""
 170
 171     # read the zk configuration from the app's config
 172     zk_hostlist = config.get("DEFAULT", "zookeeper_connection_string")
 173     username = config.get("DEFAULT", "zookeeper_username")
 174     password = config.get("DEFAULT", "zookeeper_password")
 175
 176     client = connect_to_zookeeper(zk_hostlist, (username, password))
 177     return LiveConfig(client, LIVE_CONFIG_NODE).data
 178
 179
 180 def confirm_config(old_config, new_config):
 181     """Display the changes and confirm that we should continue."""
 182
 183     # we need to convert the new config to json and back so values like
 184     # tuples are converted to the same format as the existing config
 185     new_config = json.loads(json.dumps(new_config))
 186
 187     print ("Updates (a line starting with + is a new setting, - is removed, "
 188         "and ! is changed, with additions/removals/replacements indented):\n")
 189     print_dict_diff(old_config, new_config)
 190     answer = raw_input("\nContinue? [y|N] ")
 191     return answer.lower() == "y"
 192
 193
 194 def main():
 195     """Get and validate input from the user via CLI then write to ZK."""
 196
 197     progname = os.path.basename(sys.argv[0])
 198
 199     try:
 200         ini_file_name = sys.argv[1]
 201     except IndexError:
 202         print >> sys.stderr, "USAGE: %s INI" % progname
 203         return 1
 204
 205     try:
 206         with open(ini_file_name) as ini_file:
 207             config = parse_ini_file(ini_file)
 208     except (IOError, ConfigParser.Error), e:
 209         print >> sys.stderr, "%s: %s: %s" % (progname, ini_file_name, e)
 210         return 1
 211
 212     try:
 213         plugin_config = config.get("DEFAULT", "plugins")
 214         plugin_names = ConfigValue.tuple(plugin_config)
 215         plugins = PluginLoader(plugin_names=plugin_names)
 216         live = extract_live_config(config, plugins)
 217     except ValueError as e:
 218         print >> sys.stderr, "%s: %s" % (progname, e)
 219         return 1
 220     else:
 221         current_config = get_current_live_config(config)
 222         if not confirm_config(current_config, live):
 223             print "Oh, well, never mind then. Bye :("
 224             return 1
 225
 226     password = getpass.getpass("Password: ")
 227
 228     try:
 229         write_config_to_zookeeper(LIVE_CONFIG_NODE,
 230                                   USERNAME, password,
 231                                   config, live)
 232     except NoAuthException:
 233         print >> sys.stderr, "%s: incorrect password" % progname
 234         return 1
 235     except Exception as e:
 236         print >> sys.stderr, "%s: %s" % (progname, e)
 237         return 1
 238
 239     print "Succesfully updated live config!"
 240
 241     return 0
 242
 243
 244 if __name__ == "__main__":
 245     sys.exit(main())