New: SolrCloud support (see django-haystack#138)

This optionally adds support for SolrCloud using the Kazoo client
library.

Thanks to @upayavira

Author: acdha

.gitignore

@@ -2,9 +2,12 @@
 build
 dist/
 MANIFEST
-solr4
+solr/
+solr-app/
+logs/
 env
 env3
 pysolr.egg-info/
 .tox
 solr*.tgz
+.idea

AUTHORS

@@ -34,3 +34,4 @@ Contributors:
 * Çağatay Çallı (@faraday) for Solr 4 field update support
 * Emmanuel Leblond (@touilleMan) for fixing error handling on Python 3
 * Michał Jaworski (@swistakm) for improved Sentry-friendly logging
+* Upayavira (@upayavira) for SolrCloud support

README.rst

@@ -23,15 +23,15 @@ Features
 * `"More Like This" <http://wiki.apache.org/solr/MoreLikeThis>`_ support (if set up in Solr).
 * `Spelling correction <http://wiki.apache.org/solr/SpellCheckComponent>`_ (if set up in Solr).
 * Timeout support.
-
+* SolrCloud awareness
 
 Requirements
 ============
 
 * Python 2.7 - 3.5
 * Requests 2.0+
 * **Optional** - ``simplejson``
-
+* **Optional** - ``kazoo`` for SolrCloud mode
 
 Installation
 ============
@@ -96,6 +96,12 @@ Basic usage looks like:
     # ...or all documents.
     solr.delete(q='*:*')
 
+.. code-block:: python
+    # For SolrCloud mode, initialize your Solr like this:
+    
+    zookeeper = pysolr.Zookeeper("zkhost1:2181,zkhost2:2181,zkhost3:2181")
+    solr = pysolr.SolrCloud(zookeeper, "collection1")
+
 
 LICENSE
 =======

pysolr.py

@@ -5,12 +5,18 @@
 import datetime
 import logging
 import os
+import random
 import re
 import time
 from xml.etree import ElementTree
 
 import requests
 
+try:
+    from kazoo.client import KazooClient, KazooState
+except ImportError:
+    KazooClient = KazooState = None
+
 try:
     # Prefer simplejson, if installed.
     import simplejson as json
@@ -187,10 +193,12 @@ def is_valid_xml_char_ordinal(i):
     Char ::= #x9 | #xA | #xD | [#x20-#xD7FF] | [#xE000-#xFFFD] | [#x10000-#x10FFFF]
     """
     # conditions ordered by presumed frequency
-    return (0x20 <= i <= 0xD7FF
-            or i in (0x9, 0xA, 0xD)
-            or 0xE000 <= i <= 0xFFFD
-            or 0x10000 <= i <= 0x10FFFF)
+    return (
+        0x20 <= i <= 0xD7FF
+        or i in (0x9, 0xA, 0xD)
+        or 0xE000 <= i <= 0xFFFD
+        or 0x10000 <= i <= 0x10FFFF
+        )
 
 
 def clean_xml_string(s):
@@ -316,7 +324,8 @@ def __init__(self, url, decoder=None, timeout=60, results_cls=Results):
         self.results_cls = results_cls
 
     def __del__(self):
-        self.session.close()
+        if hasattr(self, "session"):
+            self.session.close()
 
     def _get_log(self):
         return LOG
@@ -1127,3 +1136,151 @@ def sanitize(data):
         fixed_string = fixed_string.replace(bad, good)
 
     return force_unicode(fixed_string)
+
+
+class SolrCloud(Solr):
+
+    def __init__(self, zookeeper, collection, decoder=None, timeout=60, retry_timeout=0.2, *args, **kwargs):
+        url = zookeeper.getRandomURL(collection)
+
+        super(SolrCloud, self).__init__(url, decoder=decoder, timeout=timeout, *args, **kwargs)
+
+        self.zookeeper = zookeeper
+        self.collection = collection
+        self.retry_timeout = retry_timeout
+
+    def _randomized_request(self, method, path, body, headers, files):
+        self.url = self.zookeeper.getRandomURL(self.collection)
+        LOG.debug('Using random URL: %s', self.url)
+        return Solr._send_request(self, method, path, body, headers, files)
+
+    def _send_request(self, method, path='', body=None, headers=None, files=None):
+        # FIXME: this needs to have a maximum retry counter rather than waiting endlessly
+        try:
+            return self._randomized_request(method, path, body, headers, files)
+        except requests.exceptions.RequestException:
+            LOG.warning('RequestException, retrying after %fs', self.retry_timeout, exc_info=True)
+            time.sleep(self.retry_timeout)  # give zookeeper time to notice
+            return self._randomized_request(method, path, body, headers, files)
+        except SolrError:
+            LOG.warning('SolrException, retrying after %fs', self.retry_timeout, exc_info=True)
+            time.sleep(self.retry_timeout)  # give zookeeper time to notice
+            return self._randomized_request(method, path, body, headers, files)
+
+    def _update(self, message, clean_ctrl_chars=True, commit=True, softCommit=False, waitFlush=None,
+                waitSearcher=None):
+        self.url = self.zookeeper.getLeaderURL(self.collection)
+        LOG.debug('Using random leader URL: %s', self.url)
+        return Solr._update(self, message, clean_ctrl_chars, commit, softCommit, waitFlush, waitSearcher)
+
+
+class ZooKeeper(object):
+    # Constants used by the REST API:
+    LIVE_NODES_ZKNODE = '/live_nodes'
+    ALIASES = '/aliases.json'
+    CLUSTER_STATE = '/clusterstate.json'
+    SHARDS = 'shards'
+    REPLICAS = 'replicas'
+    STATE = 'state'
+    ACTIVE = 'active'
+    LEADER = 'leader'
+    BASE_URL = 'base_url'
+    TRUE = 'true'
+    FALSE = 'false'
+    COLLECTION = 'collection'
+
+    def __init__(self, zkServerAddress, zkClientTimeout=15, zkClientConnectTimeout=15):
+        if KazooClient is None:
+            logging.error('ZooKeeper requires the `kazoo` library to be installed')
+            raise RuntimeError
+
+        self.collections = {}
+        self.liveNodes = {}
+        self.aliases = {}
+        self.state = None
+
+        self.zk = KazooClient(zkServerAddress, read_only=True)
+
+        self.zk.start()
+
+        def connectionListener(state):
+            if state == KazooState.LOST:
+                self.state = state
+            elif state == KazooState.SUSPENDED:
+                self.state = state
+        self.zk.add_listener(connectionListener)
+
+        @self.zk.DataWatch(ZooKeeper.CLUSTER_STATE)
+        def watchClusterState(data, *args, **kwargs):
+            if not data:
+                LOG.warning("No cluster state available: no collections defined?")
+            else:
+                self.collections = json.loads(data)
+                LOG.info("Updated collections")
+
+            collection_data = json.loads(data.decode('utf-8'))
+            self.collections = collection_data
+            LOG.info('Updated collections: %s', collection_data)
+
+        @self.zk.ChildrenWatch(ZooKeeper.LIVE_NODES_ZKNODE)
+        def watchLiveNodes(children):
+            self.liveNodes = children
+            LOG.info("Updated live nodes: %s", children)
+
+        @self.zk.DataWatch(ZooKeeper.ALIASES)
+        def watchAliases(data, stat):
+            if data:
+                self.aliases = json.loads(data)[ZooKeeper.COLLECTION]
+            else:
+                self.aliases = None
+            LOG.info("Updated aliases: %s", self.aliases)
+
+    def __del__(self):
+        # Avoid leaking connection handles in Kazoo's atexit handler:
+        self.zk.stop()
+        self.zk.close()
+
+    def getHosts(self, collname, only_leader=False, seen_aliases=None):
+        if self.aliases and collname in self.aliases:
+            return self.getAliasHosts(collname, only_leader, seen_aliases)
+
+        hosts = []
+        if not self.collections.has_key(collname):
+            raise SolrError("Unknown collection: %s", collname)
+        collection = self.collections[collname]
+        shards = collection[ZooKeeper.SHARDS]
+        for shardname in shards.keys():
+            shard = shards[shardname]
+            if shard[ZooKeeper.STATE] == ZooKeeper.ACTIVE:
+                replicas = shard[ZooKeeper.REPLICAS]
+                for replicaname in replicas.keys():
+                    replica = replicas[replicaname]
+
+                    if replica[ZooKeeper.STATE] == ZooKeeper.ACTIVE:
+                        if not only_leader or (replica.get(ZooKeeper.LEADER, None) == ZooKeeper.TRUE):
+                            base_url = replica[ZooKeeper.BASE_URL]
+                            if base_url not in hosts:
+                                hosts.append(base_url)
+        return hosts
+
+    def getAliasHosts(self, collname, only_leader, seen_aliases):
+        if seen_aliases:
+            if collname in seen_aliases:
+                LOG.warn("%s in circular alias definition - ignored", collname)
+                return []
+        else:
+            seen_aliases = []
+        seen_aliases.append(collname)
+        collections = self.aliases[collname].split(",")
+        hosts = []
+        for collection in collections:
+            for host in self.getHosts(collection, only_leader, seen_aliases):
+                if host not in hosts:
+                    hosts.append(host)
+        return hosts
+
+    def getRandomURL(self, collname):
+        return random.choice(self.getHosts(collname, only_leader=False)) + "/" + collname
+
+    def getLeaderURL(self, collname):
+        return random.choice(self.getHosts(collname, only_leader=True)) + "/" + collname

run-tests.py

@@ -1,49 +1,22 @@
 #!/usr/bin/env python
+# encoding: utf-8
 
 from __future__ import absolute_import, print_function, unicode_literals
 
-import subprocess
-import sys
-import time
+import unittest
 
-import requests
-
-
-def start_solr():
-    solr_proc = subprocess.Popen("./start-solr-test-server.sh",
-                                 stdout=open("test-solr.stdout.log", "wb"),
-                                 stderr=open("test-solr.stderr.log", "wb"))
-
-    solr_retries = 0
-
-    while True:
-        try:
-            r = requests.get("http://localhost:8983/solr/core0/select/?q=startup")
-            status_code = r.status_code
-        except requests.RequestException:
-            status_code = 0
-
-        if status_code == 200:
-            break
-        elif solr_retries < 60:
-            solr_retries += 1
-            print("Waiting 10 seconds for Solr to start (retry #%d)" % solr_retries, file=sys.stderr)
-            time.sleep(10)
-        else:
-            print("Solr took too long to start (#%d retries)" % solr_retries, file=sys.stderr)
-            sys.exit(1)
-
-    return solr_proc
+from tests import utils as test_utils
 
 
 def main():
-    solr_proc = start_solr()
+    test_utils.prepare()
+    test_utils.start_solr()
 
     try:
-        subprocess.check_call(['python', '-m', 'unittest', 'tests'])
+        unittest.main(module='tests', verbosity=1)
     finally:
-        solr_proc.terminate()
-        solr_proc.wait()
+        print('Tests complete; halting Solr servers…')
+        test_utils.stop_solr()
 
 if __name__ == "__main__":
     main()

setup.cfg

@@ -4,6 +4,9 @@ universal = 1
 [pep8]
 ignore = W503
 
+[pylama]
+ignore = D400,D401
+
 [isort]
 combine_as_imports = true
 default_section = THIRDPARTY

setup.py

@@ -28,5 +28,10 @@
     license='BSD',
     install_requires=[
         'requests>=2.0'
-    ]
+    ],
+    extras_require={
+        'solrcloud': [
+	    'kazoo==2.2'
+        ]
+    }
 )