← Back to branch summary

~bloodearnest/conn-check/1.5.1

« back to all changes in this revision

Viewing changes to conn_check/checks.py

  • Committer: Tarmac
  • Author(s): Wes Mason
  • Date: 2014-10-29 16:20:19 UTC
  • mfrom: (79.1.10 mongodb)
  • Revision ID: tarmac-20141029162019-g2wcrsq3s5gk32g6
[r=james-w] Add a MongoDB check type

Show diffs side-by-side

added added

removed removed

Lines of Context:
conn_check/checks.py
344
344
 
345
345
    @inlineCallbacks
346
346
    def do_connect():
347
 
        """Connect and authenticate.
348
 
        """
 
347
        """Connect and authenticate."""
349
348
        client_creator = ClientCreator(reactor, MemCacheProtocol)
350
349
        client = yield client_creator.connectTCP(host=host, port=port,
351
350
                                                 timeout=timeout)
357
356
                            sequential_check(subchecks))
358
357
 
359
358
 
 
359
def make_mongodb_check(host, port=27017, username=None, password=None,
 
360
                       database='test', timeout=10, **kwargs):
 
361
    """Return a check for MongoDB connectivity."""
 
362
 
 
363
    import txmongo
 
364
    subchecks = []
 
365
    subchecks.append(make_tcp_check(host, port, timeout=timeout))
 
366
 
 
367
    port = int(port)
 
368
 
 
369
    @inlineCallbacks
 
370
    def do_connect():
 
371
        """Try to establish a mongodb connection."""
 
372
        conn = txmongo.MongoConnection(host, port)
 
373
 
 
374
        conn.uri['options']['connectTimeoutMS'] = int(timeout*1000)
 
375
        if username:
 
376
            conn.uri['username'] = username
 
377
        if password:
 
378
            conn.uri['password'] = password
 
379
 
 
380
        # We don't start our timeout callback until now, otherwise we might
 
381
        # elapse part of our timeout period during the earlier TCP check
 
382
        reactor.callLater(timeout, timeout_handler)
 
383
 
 
384
        mongo = yield conn
 
385
        names = yield mongo[database].collection_names()
 
386
 
 
387
    def timeout_handler():
 
388
        """Manual timeout handler as txmongo timeout args above don't work in
 
389
        many situations."""
 
390
        if 'deferred' in do_connect.func_dict:
 
391
            err = ValueError("timeout connecting to mongodb")
 
392
            do_connect.func_dict['deferred'].errback(err)
 
393
 
 
394
    connect_info = "connect with auth" if any((username, password)) else "connect"
 
395
    subchecks.append(make_check(connect_info, do_connect))
 
396
    return add_check_prefix('mongodb:{}:{}'.format(host, port),
 
397
                            sequential_check(subchecks))
 
398
 
 
399
 
360
400
CHECKS = {
361
401
    'tcp': {
362
402
        'fn': make_tcp_check,
398
438
        'fn': make_memcache_check,
399
439
        'args': ['host', 'port'],
400
440
    },
 
441
    'mongodb': {
 
442
        'fn': make_mongodb_check,
 
443
        'args': ['host'],
 
444
    },
 
445
    'mongo': {
 
446
        'fn': make_mongodb_check,
 
447
        'args': ['host'],
 
448
    },
401
449
}