Exemplo n.º 1
0
class SolrClient(object):
    """Solr client class used to perform requests to a Solr instance.

    The client can use different input and output methods using Twisted
    IProducer and IConsumer.

    @param url: The URL of the Solr server.
    @param inputFactory: The input body generator. For advanced uses this
        argument is used to create custom body generators for the requests
        using Twisted's IProducer.
    """

    def __init__(self, url, inputFactory=None):
        self.url = url.rstrip('/')
        if inputFactory is None:
            self.inputFactory = SimpleXMLInputFactory()

    def _request(self, method, path, headers, bodyProducer):
        """Performs a request to a Solr client

        The request examines the response to look for wrong header status.
        Additionally, it parses the response using a ResponseConsumer and
        creates a L{SolrResponse} object which will be given to the returning
        deferred callback.

        @param method: The HTTP method of the request.
        @param path: The path of the request.
        @param headers: The headers of the request.
        @bodyProducer: The L{IBodyProducer} that generates the body of the
            request.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        result = Deferred()

        url = self.url + path
        headers.update({'User-Agent': ['txSolr']})
        headers = Headers(headers)
        _logger.debug('Requesting: [%s] %s' % (method, url))
        agent = Agent(reactor)
        d = agent.request(method, url, headers, bodyProducer)

        def responseCallback(response):
            _logger.debug('Received response from ' + url)
            try:
                if response.code == 200:
                    deliveryProtocol = ResponseConsumer(result,
                                                        JSONSolrResponse)
                    response.deliverBody(deliveryProtocol)
                else:
                    deliveryProtocol = DiscardingResponseConsumer()
                    response.deliverBody(deliveryProtocol)
                    result.errback(HTTPWrongStatus(response.code))
            except Exception as e:
                result.errback(e)

        def responseErrback(failure):
            """Unknown error from Agent.request."""
            result.errback(HTTPRequestError(failure.value))
            _logger.error(failure.value)

        d.addCallbacks(responseCallback, responseErrback)

        return result

    def _update(self, input):
        """Performs a request to the /update method of Solr.

        @param input: The L{IBodyProducer} that generates the body of the
            request.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        method = 'POST'
        path = '/update?wt=json'
        headers = {'Content-Type': [self.inputFactory.contentType]}
        _logger.debug('Updating:\n%s' % input.body)
        return self._request(method, path, headers, input)

    def _select(self, params):
        """Performs a request to the /select method of Solr.

        @param params: A C{dict} with the request parameters as C{unicode}
            used for the query.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        # force JSON response for now
        params.update(wt=u'json')

        encodedParameters = {}
        for key, value in params.iteritems():
            # Some solr params contains dots (i.e: ht.fl) We use underscores.
            key = key.replace('_', '.')
            if isinstance(value, unicode):
                value = value.encode('UTF-8')
            encodedParameters[key] = value

        query = urllib.urlencode(encodedParameters)

        if len(query) < 1024:
            method = 'GET'
            path = '/select' + '?' + query
            headers = {}
            input = None
        else:
            method = 'POST'
            path = '/select'
            headers = {
                'Content-type': [
                    'application/x-www-form-urlencoded; charset=UTF-8'
                ]
            }
            input = StringProducer(query)

        return self._request(method, path, headers, input)

    def add(self, documents, overwrite=None, commitWithin=None):
        """Add one or many documents to a Solr Instance.

        @param documents: A C{dict} or C{list} of dicts representing the
            documents. The dict's keys should be field names and the values
            field content.
        @param overwrite: Newer documents will replace previously added
            documents with the same C{uniqueKey}.
        @param commitWithin: the addition will be committed within that time.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        input = self.inputFactory.createAdd(documents, overwrite, commitWithin)
        return self._update(input)

    def delete(self, ids):
        """Delete one or many documents given the ID or IDs of the documents.

        @param ids: A C{string} or list of string representing the IDs of the
            documents that will be deleted.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """

        input = self.inputFactory.createDelete(ids)
        return self._update(input)

    def deleteByQuery(self, query):
        """Delete all documents returned by a query.

        @param query: A Solr query that returns the documents to be deleted.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """

        input = self.inputFactory.createDeleteByQuery(query)
        return self._update(input)

    def commit(self, waitFlush=None, waitSearcher=None, expungeDeletes=None):
        """Issues a commit action to Sorl.

        @param waitFlush: Server will block until index changes are flushed to
            disk.
        @param waitSearcher: Server will  block until a new searcher is opened
            and registered as the main query searchers.
        @param expungeDeletes: Merge segments with deletes away.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        input = self.inputFactory.createCommit(waitFlush, waitSearcher,
                                               expungeDeletes)
        return self._update(input)

    def rollback(self):
        """Withdraw all uncommitted changes.

        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        input = self.inputFactory.createRollback()
        return self._update(input)

    def optimize(self, waitFlush=None, waitSearcher=None, maxSegments=None):
        """Issues an optimize action to Solr.

        @param waitFlush: Server will block until index changes are flushed
            to disk.
        @param waitSearcher: Server will  block until a new searcher is opened
            and registered as the main query searcher.
        @param maxSegments: Optimizes down to at most this number of segments.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        input = self.inputFactory.createOptimize(waitFlush, waitSearcher,
                                                 maxSegments)
        return self._update(input)

    def search(self, query, **kwargs):
        """Performs a query to Solr.

        @param query: A C{unicode} query. (See Solr query syntax).
        @param *kwargs: Additional parameters for the server. For instance:
            'hl' for highlighting, 'sort' for sorting, etc. See Solr
            documentation for all available options.
        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        params = {}
        params.update(kwargs)
        params.update(q=query)
        return self._select(params)

    def ping(self):
        """Ping the server to know if it's alive.

        This will only work if the Solr server is configured for ping.

        @return: A L{Deferred} that fires with a L{SolrResponse} object.
        """
        method = 'GET'
        path = '/admin/ping?wt=json'
        headers = {}
        return self._request(method, path, headers, None)
Exemplo n.º 2
0
class SimpleXMLInputFactoryTest(unittest.TestCase):

    def setUp(self):
        self.input = SimpleXMLInputFactory()

    # TODO: Don't use private methods for this test.
    def testEncodeValue(self):
        """L{SimpleXMLInputFactory} encodes primitive values correctly."""
        value = datetime(2010, 1, 1, 0, 0, 0)
        value = self.input._encodeValue(value)
        self.assertEqual(value, '2010-01-01T00:00:00Z')

        value = date(2010, 1, 1)
        value = self.input._encodeValue(value)
        self.assertEqual(value, '2010-01-01T00:00:00Z')

        value = True
        value = self.input._encodeValue(value)
        self.assertEqual(value, 'true')

        value = 'sample str'
        value = self.input._encodeValue(value)
        self.assert_(isinstance(value, unicode))

        value = None
        value = self.input._encodeValue(value)
        self.assert_(isinstance(value, unicode))

    def testCreateAdd(self):
        """
        L{SimpleXMLInputFactory.createAdd} creates a correct body for an C{add}
        request.
        """
        document = {'id': 1, 'text': 'hello'}
        expected = ('<add><doc><field name="text">hello</field>'
                    '<field name="id">1</field></doc></add>')
        input = self.input.createAdd(document).body
        self.assertEqual(input, expected, 'Wrong input')

    def testCreateAddEncoding(self):
        """
        L{SimpleXMLInputFactory.createAdd} encodes the C{add} request as UTF-8.
        """
        document = {'id': 1, 'text': u'\U0001d1b6'}
        expected = ('<add><doc><field name="text">\xf0\x9d\x86\xb6</field>'
                    '<field name="id">1</field></doc></add>')
        input = self.input.createAdd(document).body
        self.assertEqual(input, expected, 'Wrong input')

    def testCreateAddWithCollection(self):
        """
        L{SimpleXMLInputFactory.createAdd} creates a correct body for an C{add}
        request with a sequence field.
        """
        document = {'id': 1, 'collection': [1, 2, 3]}
        expected = ('<add><doc><field name="id">1</field>'
                    '<field name="collection">1</field>'
                    '<field name="collection">2</field>'
                    '<field name="collection">3</field></doc></add>')

        input = self.input.createAdd(document).body
        self.assertEqual(input, expected, 'Wrong input')

    def testCreateAddWithWrongValues(self):
        """
        L{SimpleXMLInputFactory.createAdd} raises C{AttributeError} if one of
        the given values is not a proper C{dict} document.
        """
        self.assertRaises(AttributeError, self.input.createAdd, None)
        self.assertRaises(AttributeError, self.input.createAdd, 'string')

    def testCreateAddWithOverwrite(self):
        """
        L{SimpleXMLInputFactory.createAdd} can add an C{overwrite} option.
        """
        document = {'id': 1, 'text': 'hello'}
        expected = ('<add overwrite="true">'
                    '<doc><field name="text">hello</field>'
                    '<field name="id">1</field></doc></add>')

        input = self.input.createAdd(document, overwrite=True).body
        self.assertEqual(input, expected, 'Wrong input')

    def testCreateAddWithCommitWithin(self):
        """
        L{SimpleXMLInputFactory.createAdd} can add a C{commitWithin} option.
        """
        document = {'id': 1, 'text': 'hello'}
        expected = ('<add commitWithin="80">'
                    '<doc><field name="text">hello</field>'
                    '<field name="id">1</field></doc></add>')

        input = self.input.createAdd(document, commitWithin=80).body
        self.assertEqual(input, expected, 'Wrong input')

    def testCreateDelete(self):
        """
        L{SimpleXMLInputFactory.createDelete} creates a correct body for a
        C{delete} request.
        """
        id = 123
        expected = '<delete><id>123</id></delete>'
        self.assertEqual(self.input.createDelete(id).body, expected)

    def testCreateDeleteEncoding(self):
        """
        L{SimpleXMLInputFactory.createDelete} encodes a C{delete} request as
        UTF-8.
        """
        id = u'\U0001d1b6'
        expected = '<delete><id>\xf0\x9d\x86\xb6</id></delete>'
        self.assertEqual(self.input.createDelete(id).body, expected)

    def testCreateDeleteWithEncoding(self):
        """
        L{SimpleXMLInputFactory.createDelete} correctly escapes XML special
        characters.
        """
        id = '<hola>'
        expected = '<delete><id>&lt;hola&gt;</id></delete>'
        self.assertEqual(self.input.createDelete(id).body, expected)

    def testCreateDeleteMany(self):
        """
        L{SimpleXMLInputFactory.createDelete} correctly handles multiple
        document IDs.
        """
        id = [1, 2, 3]
        expected = '<delete><id>1</id><id>2</id><id>3</id></delete>'
        self.assertEqual(self.input.createDelete(id).body, expected)

    def testCommit(self):
        """
        L{SimpleXMLInputFactory.createCommit} creates a correct body for a
        C{commit} request.
        """
        input = self.input.createCommit().body
        expected = '<commit />'
        self.assertEqual(input, expected)

    def testCommitWaitFlush(self):
        """
        L{SimpleXMLInputFactory.createCommit} can add a C{waitFlush} parameter.
        """
        input = self.input.createCommit(waitFlush=True).body
        expected = '<commit waitFlush="true" />'
        self.assertEqual(input, expected)

    def testCommitWaitSearcher(self):
        """
        L{SimpleXMLInputFactory.createCommit} can add a C{waitSearcher}
        parameter.
        """
        input = self.input.createCommit(waitSearcher=True).body
        expected = '<commit waitSearcher="true" />'
        self.assertEqual(input, expected)

    def testCommitExpungeDeletes(self):
        """
        L{SimpleXMLInputFactory.createCommit} can add an C{expungeDeletes}
        parameter.
        """
        input = self.input.createCommit(expungeDeletes=True).body
        expected = '<commit expungeDeletes="true" />'
        self.assertEqual(input, expected)

    def testOptimize(self):
        """
        L{SimpleXMLInputFactory.createOptimize} creates a correct body for an
        C{optimize} request.
        """
        input = self.input.createOptimize().body
        expected = '<optimize />'
        self.assertEqual(input, expected)

    def testOptimizeWaitFlush(self):
        """
        L{SimpleXMLInputFactory.createOptimize} can add a C{waitFlush}
        parameter.
        """
        input = self.input.createOptimize(waitFlush=True).body
        expected = '<optimize waitFlush="true" />'
        self.assertEqual(input, expected)

    def testOptimizeWaitSearcher(self):
        """
        L{SimpleXMLInputFactory.createOptimize} can add a C{waitSearcher}
        parameter.
        """
        input = self.input.createOptimize(waitSearcher=True).body
        expected = '<optimize waitSearcher="true" />'
        self.assertEqual(input, expected)

    def testOptimizeMaxSegments(self):
        """
        L{SimpleXMLInputFactory.createCommit} can add a C{maxSegments}
        parameter.
        """
        input = self.input.createOptimize(maxSegments=2).body
        expected = '<optimize maxSegments="2" />'
        self.assertEqual(input, expected)