Runs Influxdb updates to setup new users and databases, triggered on AWS

influxdb_update.py

#!/usr/bin/env python
#
# Author: dataday (2015)
#
# Description: influxdb updates for new users and databases
#
# To remove influx data and meta
# rm -rf /var/opt/influxdb/meta/*; rm -rf /var/opt/influxdb/data/*; rm -rf /var/opt/influxdb/wal/*
# service influxdb restart
#
import json
import os
import re

def str_to_arr(comma_str):
  return [key.strip() for key in comma_str.split(',') if comma_str]

def query(query):
  os.system('/opt/influxdb/influx %s' % query)
  print re.sub(r"((with\ |[\s\-]+)?password '[\w]+')", "", query)

def get_value(data, key):
  return data.get(key)

def main (data):
  users = str_to_arr(data.get('influx-users'))
  dbs = str_to_arr(data.get('influx-dbs'))

  dbs_query = "-username {0} -password '{1}' -database {2} -execute \"{3}\""
  dbs_create_db = "create database {0}"
  dbs_create_admin = "create user {0} with password '{1}' with all privileges"
  dbs_create_user = "create user {0} with password '{1}'"
  dbs_grant_user = "grant {0} on {1} to {2}"

  if 'admin' in users:

    users.remove('admin')

    # create the admin user
    print "create admin"
    admin_pwd = get_value(data, 'influx-user-admin')
    admin_cmd = dbs_query.format('admin', admin_pwd, 'mydb', '{0}')
    query(dbs_query.format('root', 'root', 'mydb', dbs_create_admin.format('admin', admin_pwd)))

    # create a grafana user
    if 'grafana' in users:
      print "create user grafana"
      grafana_pwd = get_value(data, 'influx-user-grafana')
      query(admin_cmd.format(dbs_create_user.format('grafana', grafana_pwd)))
    else:
      print "Grafana hasn't been set-up"

    # create specified databases
    for db in dbs:
      print "create database: %s" % db
      query(dbs_query.format('admin', admin_pwd, 'mydb', dbs_create_db.format(db)))
      if 'grafana' in users:
        # grant database read permission to grafana
        print "update database {0} with user: grafana".format(db)
        query(admin_cmd.format(dbs_grant_user.format('read', db, 'grafana')))

    # create users associated to a database
    # e.g., user = foo, database = foo
    # initial creation: database and user need to have matching names
    for user in users:
      if user in dbs:
        print "add database {0} with user: {0}".format(user)
        user_pwd = get_value(data, 'influx-user-{0}'.format(user))
        query(admin_cmd.format(dbs_create_user.format(user, user_pwd)))
        query(admin_cmd.format(dbs_grant_user.format('all', user, user)))

  else:
    print "Admin is needed in order to update influx"

if __name__ == "__main__":

  data = json.load(open('/etc/config.json'))
  main(data['secure_configuration'])

''' SAMPLECONFIG
[
  {
    "key": "influx-user-admin",
    "value": "PASSWORD1",
    "secure": true
  },
  {
    "key": "influx-user-db1user",
    "value": "PASSWORD1",
    "secure": true
  },
  {
    "key": "influx-user-db2user",
    "value": "PASSWORD1",
    "secure": true
  },
  {
    "key": "influx-user-grafana",
    "value": "PASSWORD1",
    "secure": true
  },
  {
    "key": "influx-user-gatling",
    "value": "PASSWORD1",
    "secure": true
  },
  {
    "key": "influx-users",
    "value": "admin, db1user, db2user, grafana, gatling",
    "secure": true
  },
  {
    "key": "influx-dbs",
    "value": "db1, db2, gatling",
    "secure": true
  }
]
'''