def test_get_collection_invalid_names():
"""Test get_collection() with invalid names."""
with pytest.raises(pymongo.errors.InvalidName):
collection = mongo.get_collection('', 'collection_name')
with pytest.raises(pymongo.errors.InvalidName):
collection = mongo.get_collection('db_name', '')
def jobs_collection(self):
if not self.__jobs:
self.__jobs = common_mongo.get_collection(
self.database,
self.jobs_collection_name,
self.connection)
return self.__jobs
def tasks_collection(self):
if not self.__tasks:
self.__tasks = common_mongo.get_collection(
self.database,
self.tasks_collection_name,
self.connection
)
return self.__tasks
def __init__(self, database_servers, replica_set, database, collection):
connection = get_connection(server_or_servers=database_servers,
replica_set=replica_set)
self.collection = get_collection(database,
collection,
connection=connection)
self.db = InventoryDBController(self.collection)
super(InventoryController, self).__init__()
def get_jobs_collection(connection=None):
jobs_db = db_configuration.get('jobs_mongo_db', 'test')
jobs_collection = db_configuration.get('jobs_mongo_collection', 'rpc_jobs')
jobs_servers = db_configuration.get('jobs_mongo_servers', 'localhost')
replica_set = db_configuration.get('jobs_replica_set')
return get_collection(jobs_db,
jobs_collection,
connection=connection,
server_or_servers=jobs_servers,
replica_set=replica_set)
def get_tasks_collection(connection=None):
tasks_db = db_configuration.get('tasks_mongo_db', 'test')
tasks_collection = db_configuration.get('tasks_mongo_collection',
'rpc_tasks')
tasks_servers = db_configuration.get('tasks_mongo_servers', 'localhost')
replica_set = db_configuration.get('tasks_replica_set')
return get_collection(tasks_db,
tasks_collection,
connection=connection,
server_or_servers=tasks_servers,
replica_set=replica_set)
def main():
config = options()
logging.basicConfig(level=logging.getLevelName(config.log_level),
format=config.log_format)
db_connection = get_connection(config.log_service.db.servers,
config.log_service.db.replica_name)
collection = get_collection(config.log_service.db.name,
config.log_service.db.collection,
db_connection)
agent_log_service = AgentLogService(config.log_service.bind_address,
collection)
LOG.info('Starting logging service on {}'.format(
config.log_service.bind_address))
agent_log_service.start()
def __init__(self, inventory_configuration):
self.inventory_configuration = inventory_configuration
self.db_configuration = self.inventory_configuration.get(
'inventory', {}).get('db', {})
log.debug('DB Configuration: %s' % self.db_configuration)
db_name = self.db_configuration.get('name')
if not db_name:
log.warning('DB name is not specified, using test')
db_name = 'test'
connection = get_connection(
server_or_servers=self.db_configuration.get(
'servers', ['localhost']),
replica_set=self.db_configuration.get('replica_set'))
self.collection = get_collection(
db_name, # Perhaps we should raise an exception
self.db_configuration.get('collection', 'inventory'),
connection=connection)
self.db = InventoryDBController(self.collection)
super(InventoryController, self).__init__()
def main():
config = parse_options()
logging.basicConfig(level=logging.getLevelName(config.logging.level),
format=config.logging.format)
loop = zmq.asyncio.ZMQEventLoop()
loop.set_debug(config.asyncio_debug)
asyncio.set_event_loop(loop)
async_connection = get_async_connection(
config.inventory.database.servers,
config.inventory.database.replica_name,
username=config.inventory.database.username,
password=config.inventory.database.password
)
inventory_collection = get_collection(
config.inventory.database.name,
config.inventory.database.collection,
connection=async_connection
)
s = InventoryServer(bind_address=config.inventory.bind_address,
inventory_collection=inventory_collection)
try:
loop.run_until_complete(s.start())
except KeyboardInterrupt:
log.info('Stopping service')
s.kill()
finally:
pending = asyncio.Task.all_tasks(loop=loop)
log.debug('Waiting on {} pending tasks'.format(len(pending)))
loop.run_until_complete(asyncio.gather(*pending))
log.debug('Shutting down event loop')
loop.close()
socket.send(msgpack.packb(_payload))
if not socket.poll(timeout):
log.debug('Ping timeout: %s' % zurl)
socket.close()
return False
reply = socket.recv()
log.debug("%s : %s" % (zurl, msgpack.unpackb(reply, encoding='utf-8')))
socket.close()
return True
if __name__ == '__main__':
from mercury.rpc.db import ActiveInventoryDBController
from mercury.rpc.configuration import rpc_configuration
from mercury.common.mongo import get_collection
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s : %(levelname)s - %(name)s - %(message)s')
logging.getLogger('mercury.rpc.ping').setLevel(logging.DEBUG)
db_configuration = rpc_configuration.get('db', {})
collection = get_collection(
db_configuration.get('rpc_mongo_db', 'test'),
db_configuration.get('rpc_mongo_collection', 'rpc'),
server_or_servers=db_configuration.get('rpc_mongo_servers',
'localhost'),
replica_set=db_configuration.get('replica_set'))
_db_controller = ActiveInventoryDBController(collection=collection)
def test_get_collection():
"""Test get_collection()"""
collection = mongo.get_collection('db_name', 'collection_name')
assert isinstance(collection, pymongo.collection.Collection)
assert collection.full_name == 'db_name.collection_name'
# if not self.validate_message(message):
# raise MercuryGeneralException('Logging controller recieved invalid message')
message.update({'time_created': time.time()})
self.set_job_info_from_thread(message)
LOG.debug(message)
self.log_collection.insert(message)
return {'message': 'ok'}
if __name__ == '__main__':
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s : %(levelname)s - %(name)s - %(message)s')
configuration = get_configuration(__defaults['configuration_file'])
db_connection = get_connection(
configuration['db']['servers'],
replica_set=configuration['db']['replica_set'])
_log_collection = get_collection(configuration['db']['database'],
configuration['db']['collection'],
db_connection)
_bind_address = configuration['service']['url']
agent_log_service = AgentLogService(_bind_address, _log_collection)
LOG.info('Starting logging service on {}'.format(_bind_address))
agent_log_service.start()
