Utility script for data stored in Mongodb by Presage2

presage2_mongo.py

#!/usr/bin/python
import sys
import argparse
from pymongo import Connection

## Add optional arguments for the mongodb connection
def add_mongo_args(parser):
	parser.add_argument('--host', default='localhost', help='mongodb host')
	parser.add_argument('--port', type=int, default=27017, help='mongodb port')
	parser.add_argument('--db', default='presage', help='mongodb database name')

## Connect to the mongodb specified in args
def connect_mongodb(args):
	connection = Connection(args.host, args.port)
	db = connection[args.db]
	print 'Connected to mongodb at {0}:{1}/{2}'.format(args.host, args.port, args.db)
	return db

## Delete simulations specified by args
def delete_sims(args):
	db = connect_mongodb(args)
	simulations = db.simulations
	environment = db.environmentState
	environment_time = db.environmentStateTime
	agents = db.agents
	agent_state = db.agentstate
	for i in args.simIds:
		print "Deleting simulation ID: {}".format(i)
		environment.remove({"simId": i})
		environment_time.remove({"simId": i})
		print " - Environment state deleted"
		agent_count = 0
		for agent in agents.find({"simID": i}):
			agent_state.remove({"aid": agent["aid"]})
			agent_count = agent_count + 1
		agents.remove({"simID": i})
		print " - {} agents deleted".format(agent_count)
		simulations.remove({"_id": i})
		print " - purge complete"

# arg parse config
parser = argparse.ArgumentParser(description='')
add_mongo_args(parser)

subparsers = parser.add_subparsers(help='command help')

delete_parser = subparsers.add_parser('delete')
delete_parser.add_argument('simIds', nargs='+', type=int, metavar='ID', help='simulation[s] to delete')
add_mongo_args(delete_parser)
delete_parser.set_defaults(func=delete_sims)

args = parser.parse_args()
args.func(args)