From 0b09f25b286505e2dac19abc1f663869fdcaaf59 Mon Sep 17 00:00:00 2001
From: Neil Williams <neil@reddit.com>
Date: Mon, 30 Jul 2012 15:39:16 -0700
Subject: [PATCH] Add basic support for ZooKeeper via Kazoo.

At this point, we can connect and pass authentication credentials!
---
 r2/example.ini           |  6 +++++
 r2/r2/lib/app_globals.py |  9 ++++++++
 r2/r2/lib/zookeeper.py   | 47 ++++++++++++++++++++++++++++++++++++++++
 r2/setup.py              |  1 +
 4 files changed, 63 insertions(+)
 create mode 100644 r2/r2/lib/zookeeper.py

diff --git a/r2/example.ini b/r2/example.ini
index e7c0db3e0..efdcf41ce 100644
--- a/r2/example.ini
+++ b/r2/example.ini
@@ -183,6 +183,12 @@ amqp_user = reddit
 amqp_pass = reddit
 amqp_virtual_host = /
 
+## -- zookeeper --
+# optional at the moment
+zookeeper_connection_string =
+zookeeper_username =
+zookeeper_password =
+
 ## -- database setup --
 # list of all databases named in the subsequent table
 databases = main, comment, vote, email, authorize, award, hc, traffic
diff --git a/r2/r2/lib/app_globals.py b/r2/r2/lib/app_globals.py
index 372e29927..5b6a1ec9f 100755
--- a/r2/r2/lib/app_globals.py
+++ b/r2/r2/lib/app_globals.py
@@ -231,6 +231,15 @@ class Globals(object):
 
         self.cache_chains = {}
 
+        # for now, zookeeper will be an optional part of the stack.
+        zk_hosts = self.config.get("zookeeper_connection_string")
+        if zk_hosts:
+            from r2.lib.zookeeper import connect_to_zookeeper
+            zk_username = self.config["zookeeper_username"]
+            zk_password = self.config["zookeeper_password"]
+            self.zookeeper = connect_to_zookeeper(zk_hosts, (zk_username,
+                                                             zk_password))
+
         self.lock_cache = CMemcache(self.lockcaches, num_clients=num_mc_clients)
         self.make_lock = make_lock_factory(self.lock_cache)
 
diff --git a/r2/r2/lib/zookeeper.py b/r2/r2/lib/zookeeper.py
new file mode 100644
index 000000000..cf8c522b7
--- /dev/null
+++ b/r2/r2/lib/zookeeper.py
@@ -0,0 +1,47 @@
+# The contents of this file are subject to the Common Public Attribution
+# License Version 1.0. (the "License"); you may not use this file except in
+# compliance with the License. You may obtain a copy of the License at
+# http://code.reddit.com/LICENSE. The License is based on the Mozilla Public
+# License Version 1.1, but Sections 14 and 15 have been added to cover use of
+# software over a computer network and provide for limited attribution for the
+# Original Developer. In addition, Exhibit A has been modified to be consistent
+# with Exhibit B.
+#
+# Software distributed under the License is distributed on an "AS IS" basis,
+# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
+# the specific language governing rights and limitations under the License.
+#
+# The Original Code is reddit.
+#
+# The Original Developer is the Initial Developer.  The Initial Developer of
+# the Original Code is reddit Inc.
+#
+# All portions of the code written by reddit are Copyright (c) 2006-2012 reddit
+# Inc. All Rights Reserved.
+###############################################################################
+
+import functools
+
+from kazoo.client import KazooClient
+from kazoo.security import make_digest_acl
+
+
+def connect_to_zookeeper(hostlist, credentials):
+    """Create a connection to the ZooKeeper ensemble.
+
+    If authentication credentials are provided (as a two-tuple: username,
+    password), we will ensure that they are provided to the server whenever we
+    establish a connection.
+
+    """
+
+    client = KazooClient(hostlist,
+                         timeout=5,
+                         max_retries=3)
+
+    # convenient helper function for making credentials
+    client.make_acl = functools.partial(make_digest_acl, *credentials)
+
+    client.connect()
+    client.add_auth("digest", ":".join(credentials))
+    return client
diff --git a/r2/setup.py b/r2/setup.py
index bea4da608..c57db7b24 100644
--- a/r2/setup.py
+++ b/r2/setup.py
@@ -96,6 +96,7 @@ setup(
         "snudown",
         "l2cs",
         "lxml",
+        "kazoo",
     ],
     dependency_links=[
         "https://github.com/downloads/reddit/pylibmc/pylibmc-1.2.1-dev.tar.gz#egg=pylibmc-1.2.1-dev",