start-postgres

#!/usr/bin/python

from jinja2 import Environment, FileSystemLoader
from subprocess import check_call

import grp
import pwd
import os
import shutil

def chmod(path, mode):
  for root, dirs, files in os.walk(path):
    for dir in dirs:
      os.chmod(os.path.join(root, dir), mode)
    for file in files:
      os.chmod(os.path.join(root, file), mode)
  os.chmod(path, mode)

def chown(path, user, group):
  uid = pwd.getpwnam(user).pw_uid
  gid = grp.getgrnam(group).gr_gid
  for root, dirs, files in os.walk(path):
    for dir in dirs:
      os.chown(os.path.join(root, dir), uid, gid)
    for file in files:
      os.chown(os.path.join(root, file), uid, gid)
  os.chown(path, uid, gid)

def empty(directory):
	for file in os.listdir(directory):
		path = os.path.join(directory, file)
		try:
			if os.path.isfile(path):
				os.unlink(path)
			elif os.path.isdir(path):
				shutil.rmtree(path)
		except Exception as exception:
			print exception

def pg_get_configuration():
	configuration = { }
	# Get the initialization flag.
	value = os.getenv('PG_INIT', 'True')
	configuration.update({ 'initialize': value == 'True' })
	# Get the deployment type.
	value = os.getenv('PG_DEPLOYMENT_TYPE', 'standalone')
	configuration.update({ 'deployment_type': value })
	# Get the container role.
	value = os.getenv('PG_ROLE', 'master')
	configuration.update({ 'role': value })
	# Get the postgres user password.
	value = os.getenv('PG_PASSWORD', 'bettervoice')
	configuration.update({ 'password': value })
	# Get the IP addresses that are allowed to connect to the db.
	value = os.getenv('PG_ALLOWED_ADDRESSES', '0.0.0.0/0').split(',')
	configuration.update({ 'allowed_addresses': value })
	# Get the database data directory.
	value = os.getenv('PG_DATA_DIRECTORY', '/var/lib/postgresql/9.3/main')
	configuration.update({ 'data_directory': value })
	# Get the maximum number of connections allowed to the db.
	value = os.getenv('PG_MAX_CONNECTIONS', 100)
	configuration.update({ 'max_connections': value })
	# Get the port number for the db to listen on.
	value = os.getenv('PG_PORT', 5432)
	configuration.update({ 'port': value })
	# Get the database share buffers size.
	value = os.getenv('PG_SHARED_BUFFERS', '128')
	configuration.update({ 'shared_buffers': value })
	# Get the database write ahead log level.
	value = os.getenv('PG_WAL_LEVEL', 'minimal')
	configuration.update({ 'wal_level': value })
	# Get the number of database checkpoint segments.
	value = os.getenv('PG_CHECKPOINT_SEGMENTS', 32)
	configuration.update({ 'checkpoint_segments': value })
	# Get the number of write ahead log senders.
	value = os.getenv('PG_MAX_WAL_SENDERS', 3)
	configuration.update({ 'max_wal_senders': value })
	# Get the number of write ahead log segments to keep around for replication.
	value = os.getenv('PG_WAL_KEEP_SEGMENTS', 8)
	configuration.update({ 'wal_keep_segments': value })
	# Get the master host for replication.
	value = os.getenv('PG_MASTER_HOST', None)
	configuration.update({ 'master_host': value })
	# Get the master host for replication.
	value = os.getenv('PG_MASTER_HOST_PORT', 5432)
	configuration.update({ 'master_host_port': value })
	# Get the IP addresses for slaves.
	value = os.getenv('PG_SLAVES', '0.0.0.0/0').split(',')
	configuration.update({ 'slaves': value })
	# Get the slave heartbeat interval.
	value = os.getenv('HEARTBEAT_INTERVAL', 5)
	configuration.update({ 'hearbeat_interval': value })
	# Get the slave heartbeat flag.
	value = os.getenv('HEARTBEAT_ENABLED', 'True')
	configuration.update({ 'heartbeat_enabled': value == 'True' })
	# Get the wal-e enabled flag.
	value = os.getenv('PG_WAL_E_ENABLED', 'False')
	configuration.update({ 'wale_enabled': value == 'True' })
	# Get the AWS secret key.
	value = os.getenv('AWS_SECRET_ACCESS_KEY', None)
	configuration.update({ 'aws_secret_access_key': value })
	# Get the AWS access key.
	value = os.getenv('AWS_ACCESS_KEY_ID', None)
	configuration.update({ 'aws_access_key_id': value })
	# Get the wal-e S3 prefix.
	value = os.getenv('WALE_S3_PREFIX', None)
	configuration.update({ 'wale_s3_prefix': value })
	return configuration

def pg_base_backup(data_directory):
	run('sudo', '-u', 'postgres', 'envdir', '/etc/wal-e.d/env', 'wal-e', 'backup-push', data_directory)

def pg_change_data_directory(data_directory):
	run('pg_dropcluster', '--stop', '9.3', 'main')
	run('pg_createcluster', '-d', data_directory, '9.3', 'main')
	chown(data_directory, 'postgres', 'postgres')

def pg_run_query(query, database = None):
	if database:
		run('sudo', '-u', 'postgres', 'psql', '-U', 'postgres', '-d', database, '-c', query)
	else:
		run('sudo', '-u', 'postgres', 'psql', '-U', 'postgres', '-c', query)

def pg_sync(host, data_directory):
	empty(data_directory)
	run('pg_basebackup', '-h', host, '-D', data_directory, '-U', 'postgres', '-v', '-x')
	chown(data_directory, 'postgres', 'postgres')

def pg_set_password(password):
	pg_run_query('alter user postgres with password \'' + password + '\';', database = 'postgres')

def run(app, *args):
	check_call([app] + list(args))

def start(service):
	check_call(['service', service, 'start'])

def wale_init(configuration):
	wale_configuration_dir = '/etc/wal-e.d/env'
	if not os.path.exists(wale_configuration_dir):
		os.makedirs(wale_configuration_dir)
	write(configuration['aws_secret_access_key'], wale_configuration_dir + '/AWS_SECRET_ACCESS_KEY')
	write(configuration['aws_access_key_id'], wale_configuration_dir + '/AWS_ACCESS_KEY_ID')
	write(configuration['wale_s3_prefix'], wale_configuration_dir + '/WALE_S3_PREFIX')
	chown(wale_configuration_dir, 'root', 'postgres')

def write(template, path):
	with open(path, "wb") as output:
		output.write(template)

if __name__ == "__main__":
	# Initialize Jinja2
	loader = FileSystemLoader('/usr/share/postgresql/9.3')
	templates = Environment(loader = loader)
	# Load the configuration into a dictionary.
	configuration = pg_get_configuration()
	if not configuration['initialize']:
		start('postgresql')
		if configuration['deployment_type'] == 'replicated' and \
			 configuration['role'] == 'slave' and \
			 configuration['heartbeat_enabled']:
			run('heartbeat')
	else:
		# Initialize the data directory.
		pg_change_data_directory(configuration['data_directory'])
		# Configure the slave.
		if configuration['deployment_type'] == 'replicated' and \
			 configuration['role'] == 'slave':
			# Copy data from the master and create the recovery file.
			pg_sync(configuration['master_host'], configuration['data_directory'])
			recovery = templates.get_template('recovery.conf.template') \
													.render(configuration)
			write(recovery, configuration['data_directory'] + '/recovery.conf')
			chown(configuration['data_directory'] + '/recovery.conf', 'postgres', 'postgres')
		# Write pg_hba.conf
		pg_hba = templates.get_template('pg_hba.conf.template') \
											.render(configuration)
		write(pg_hba, '/etc/postgresql/9.3/main/pg_hba.conf')
		# Write postgresql.conf
		postgresql = templates.get_template('postgresql.conf.template') \
													.render(configuration)
		write(postgresql, '/etc/postgresql/9.3/main/postgresql.conf')
		# Start the container.
		start('postgresql')
		# Set the default password.
		if configuration['deployment_type'] == 'standalone' or \
		   configuration['deployment_type'] == 'replicated' and \
		   configuration['role'] == 'master':
			pg_set_password(configuration['password'])
		# Post startup initialization.
		if configuration['wale_enabled']:
			wale_init(configuration)
			pg_base_backup(configuration['data_directory'])
		if configuration['deployment_type'] == 'replicated' and \
			 configuration['role'] == 'slave':
			heartbeat = templates.get_template('heartbeat.template') \
													 .render(configuration)
			write(heartbeat, '/usr/bin/heartbeat')
			chmod('/usr/bin/heartbeat', 0755)
			if configuration['heartbeat_enabled']:
				run('heartbeat')
	# Print the log to stdout.
	try:
		run('tail', '-F', '/var/log/postgresql/postgresql-9.3-main.log')
	except KeyboardInterrupt:
		print 'Exiting...'