|
#!/usr/bin/env python |
|
# coding: utf-8 |
|
|
|
import sys, os, argparse, logging |
|
import requests, json |
|
from faker import Factory |
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--db', type=str, required=True) |
|
parser.add_argument('--ddoc', type=str, required=True) |
|
parser.add_argument('--count', type=int, default=100) |
|
parser.add_argument('--user', type=str) |
|
parser.add_argument('--password', type=str) |
|
return parser.parse_args() |
|
|
|
def init_logging(): |
|
requests_log = logging.getLogger('requests.packages.urllib3.connectionpool') |
|
logging.basicConfig( |
|
format='%(asctime)s %(message)s', |
|
datefmt="%H:%M:%S", |
|
level=logging.INFO) |
|
requests_log.disabled = True |
|
|
|
def init_config(user=None, password=None): |
|
global base_url |
|
global headers |
|
global auth |
|
base_url = 'http://localhost:5984' |
|
headers = {'content-type': 'application/json'} |
|
if user is not None and password is not None: |
|
auth = (user, password) |
|
else: |
|
auth = () |
|
|
|
def init_faker(): |
|
global fake |
|
fake = Factory.create() |
|
|
|
def create_db(db_name): |
|
url = '{base_url}/{db_name}'.format(base_url=base_url, db_name=db_name) |
|
requests.delete(url, auth=auth, headers=headers) |
|
resp = requests.put(url, auth=auth, headers=headers) |
|
if resp.status_code != 201: |
|
raise ValueError(resp.json()) |
|
logging.info('Created database {db_name}'.format(db_name=db_name)) |
|
return db_name |
|
|
|
def generate_data(count): |
|
data = {} |
|
for _ in xrange(1, count): |
|
state = fake.state() |
|
city = fake.city() |
|
if state not in data: |
|
data[state] = [] |
|
data[state].append(city) |
|
return data |
|
|
|
def propagate_database(db_name, data): |
|
url = '{base_url}/{db_name}'.format(base_url=base_url, db_name=db_name) |
|
for state, cities in data.iteritems(): |
|
logging.info('Walking state {}'.format(state)) |
|
links = [] |
|
for city in cities: |
|
doc = json.dumps({ |
|
'name': city, |
|
'type': 'city', |
|
'latitude': str(fake.latitude()), |
|
'longitude': str(fake.longitude()), |
|
'postcode': fake.postcode(), |
|
'population': fake.random_int(2000, 600000), |
|
'motto': fake.sentence(), |
|
'description': fake.paragraph(), |
|
'mayor': fake.name(), |
|
'url': fake.uri() |
|
}) |
|
resp = requests.post(url, auth=auth, headers=headers, data=doc) |
|
r = resp.json() |
|
logging.info(' Added city {} ({})'.format(city, r['id'])) |
|
link = {'name': city, 'id': r['id']} |
|
links.append(link) |
|
doc = json.dumps({'name': state, 'type': 'state', 'links': links}) |
|
resp = requests.post(url, auth=auth, headers=headers, data=doc) |
|
r = resp.json() |
|
logging.info('Added state {} ({})'.format(state, r['id'])) |
|
|
|
def create_views(db_name, ddoc_name): |
|
url = '{base_url}/{db_name}/_design/{ddoc_name}'.format(base_url=base_url, |
|
db_name=db_name, ddoc_name=ddoc_name) |
|
doc = json.dumps({ |
|
'language': 'javascript', |
|
'views': { |
|
'states': { |
|
'map': ''' |
|
function (doc) { |
|
if (doc.type == 'state') { |
|
emit(doc.name, doc.links.length); |
|
} |
|
} |
|
''' |
|
}, |
|
'cities': { |
|
'map': ''' |
|
function (doc) { |
|
if (doc.type == 'state') { |
|
if (doc.links) { |
|
for (var i in doc.links) { |
|
emit([doc.name, doc.links[i].name], {_id: doc.links[i].id}); |
|
} |
|
} |
|
} |
|
} |
|
''' |
|
} |
|
} |
|
}) |
|
resp = requests.put(url, auth=auth, headers=headers, data=doc) |
|
r = resp.json() |
|
logging.info('Added ddoc {} (success: {})'.format(ddoc_name, r['ok'])) |
|
|
|
def main(): |
|
args = parse_args() |
|
init_logging() |
|
init_config(args.user, args.password) |
|
init_faker() |
|
logging.info('Creating database {}'.format(args.db)) |
|
create_db(args.db) |
|
data = generate_data(args.count) |
|
propagate_database(args.db, data) |
|
create_views(args.db, args.ddoc) |
|
sys.exit(0) |
|
|
|
if __name__ == '__main__': |
|
main() |