def handle_dropdb(self, superuser_pg_url, database_name):
"""Handles dropping the database
If the database is not a test database, this will prompt the user to make sure the user is
ok with dropping the database.
:arg string superuser_pg_url: the postgres url for the superuser user
:arg string database_name: the name of the database to drop
:returns: True if everything is good, False if the user said "don't drop this"
"""
self.config.logger.info('drop database section with %s',
superuser_pg_url)
with PostgreSQLAlchemyManager(superuser_pg_url,
self.config.logger,
autocommit=False) as db:
if 'test' not in database_name and not self.config.force:
confirm = raw_input('drop database %s [y/N]: ' % database_name)
if not confirm == "y":
self.config.logger.warning('NOT dropping table')
return False
db.drop_database(database_name)
db.commit()
return True
def main(self):
database_name = self.config.database_name
if not database_name:
self.config.logger.error(
'"database_name" cannot be an empty string'
)
return 1
# superuser credentials for overall database
superuser_pg_url = self.create_connection_url(
'postgres',
self.config.database_superusername,
self.config.database_superuserpassword
)
# superuser credentials for working database
superuser_normaldb_pg_url = self.create_connection_url(
database_name,
self.config.database_superusername,
self.config.database_superuserpassword
)
# normal user credentials
normal_user_pg_url = self.create_connection_url(
database_name,
self.config.database_username,
self.config.database_password
)
# ensure that if on Heroku the the normal_user_pg_url and the
# superuser_pg_url are the same
if self.config.on_heroku and (normal_user_pg_url != superuser_pg_url):
self.config.logger.error(
'there is no superuser (%s) when using Heroku',
self.config.database_superusername
)
return 1
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# table logging section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if self.config.unlogged:
@compiles(CreateTable)
def create_table(element, compiler, **kw):
text = compiler.visit_create_table(element, **kw)
text = re.sub("^\sCREATE(.*TABLE)",
lambda m: "CREATE UNLOGGED %s" %
m.group(1), text)
return text
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Postgres version check section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
self.config.logger.info(
'Postgres version check section with %s',
superuser_pg_url
)
with PostgreSQLAlchemyManager(
superuser_pg_url,
self.config.logger,
autocommit=False
) as db:
if not db.min_ver_check(90200):
self.config.logger.error(
'unrecognized PostgreSQL version: %s',
db.version_string()
)
self.config.logger.error('Only 9.2+ is supported at this time')
return 1
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# drop database section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# We can only do the following if the DB is not Heroku
# XXX Might add the heroku commands for resetting a DB here
if self.config.dropdb and not self.config.on_heroku:
self.config.logger.info(
'drop database section with %s',
superuser_pg_url
)
with PostgreSQLAlchemyManager(
superuser_pg_url,
self.config.logger,
autocommit=False
) as db:
if 'test' not in database_name and not self.config.force:
confirm = raw_input(
'drop database %s [y/N]: ' % database_name)
if not confirm == "y":
self.config.logger.warn('NOT dropping table')
return 2
db.drop_database(database_name)
db.commit()
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# create database section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if self.config.createdb:
self.config.logger.info(
'create database section with %s',
superuser_pg_url
)
with PostgreSQLAlchemyManager(
superuser_pg_url,
self.config.logger,
autocommit=False
) as db:
db.create_database(database_name)
if self.config.no_roles:
self.config.logger.info("Skipping role creation")
else:
db.create_roles(self.config)
db.commit()
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# database extensions section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
self.config.logger.info(
'database extensions section with %s',
superuser_normaldb_pg_url
)
with PostgreSQLAlchemyManager(
superuser_normaldb_pg_url,
self.config.logger,
autocommit=False,
on_heroku=self.config.on_heroku
) as db:
db.setup_extensions()
db.grant_public_schema_ownership(self.config.database_username)
db.commit()
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# database schema section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if self.config.no_schema:
self.config.logger.info("not adding a schema")
return 0
alembic_cfg = Config(self.config.alembic_config)
alembic_cfg.set_main_option('sqlalchemy.url', normal_user_pg_url)
self.config.logger.info(
'database schema section with %s',
normal_user_pg_url
)
with PostgreSQLAlchemyManager(
normal_user_pg_url,
self.config.logger,
autocommit=False,
on_heroku=self.config.on_heroku
) as db:
# Order matters below
db.turn_function_body_checks_off()
db.load_raw_sql('types')
db.load_raw_sql('procs')
# We need to commit to make a type visible for table creation
db.commit()
db.create_tables()
db.load_raw_sql('views')
db.commit()
if not self.config.get('no_staticdata'):
self.import_staticdata(db)
if self.config['fakedata']:
self.generate_fakedata(db, self.config['fakedata_days'])
db.commit()
command.stamp(alembic_cfg, "heads")
db.session.close()
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# database owner section
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
self.config.logger.info(
'database extensions section with %s',
superuser_normaldb_pg_url
)
with PostgreSQLAlchemyManager(
superuser_normaldb_pg_url,
self.config.logger,
autocommit=False,
on_heroku=self.config.on_heroku
) as db:
db.set_table_owner(self.config.database_username)
db.set_default_owner(database_name, self.config.database_username)
db.set_grants(self.config) # config has user lists
return 0
def main(self):
database_name = self.config.database_name
if not database_name:
self.config.logger.error(
'"database_name" cannot be an empty string')
return 1
# superuser credentials for overall database
superuser_pg_url = self.create_connection_url(
'postgres', self.config.database_superusername,
self.config.database_superuserpassword)
# superuser credentials for working database
superuser_normaldb_pg_url = self.create_connection_url(
database_name, self.config.database_superusername,
self.config.database_superuserpassword)
# normal user credentials
normal_user_pg_url = self.create_connection_url(
database_name, self.config.database_username,
self.config.database_password)
# table logging section
if self.config.unlogged:
@compiles(CreateTable)
def create_table(element, compiler, **kw):
text = compiler.visit_create_table(element, **kw)
text = re.sub("^\sCREATE(.*TABLE)",
lambda m: "CREATE UNLOGGED %s" % m.group(1),
text)
return text
# Postgres version check section
self.config.logger.info('Postgres version check section with %s',
superuser_pg_url)
with PostgreSQLAlchemyManager(superuser_pg_url,
self.config.logger,
autocommit=False) as db:
if not db.min_ver_check(90200):
self.config.logger.error('unrecognized PostgreSQL version: %s',
db.version_string())
self.config.logger.error('Only 9.2+ is supported at this time')
return 1
# At this point, we're done with setup and we can perform actions requested by the user
if self.config.dropdb:
# If this doesn't return True (aka everything worked fine), then return exit code 2
if not self.handle_dropdb(superuser_pg_url, database_name):
return 2
if self.config.createdb:
# If this doesn't return True (aka everything worked fine), then return exit code 1
if not self.handle_createdb(superuser_normaldb_pg_url,
superuser_pg_url, normal_user_pg_url,
database_name):
return 1
return 0
def handle_createdb(self, superuser_normaldb_pg_url, superuser_pg_url,
normal_user_pg_url, database_name):
"""Handles creating the database and populating it with critical stuff
:arg string superuser_normaldb_pg_url: super creds for working db
:arg string superuser_pg_url: superuser creds for overall db
:arg string normal_user_pg_url: normal user creds
:arg string database_name: the name of the database to create
:returns: True if everything worked
"""
self.config.logger.info('create database section with %s',
superuser_pg_url)
with PostgreSQLAlchemyManager(superuser_pg_url,
self.config.logger,
autocommit=False) as db:
db.create_database(database_name)
if self.config.no_roles:
self.config.logger.info("Skipping role creation")
else:
db.create_roles(self.config)
db.commit()
# database extensions section
self.config.logger.info('database extensions section with %s',
superuser_normaldb_pg_url)
with PostgreSQLAlchemyManager(superuser_normaldb_pg_url,
self.config.logger,
autocommit=False) as db:
db.setup_extensions()
db.grant_public_schema_ownership(self.config.database_username)
db.commit()
# database schema section
if self.config.no_schema:
self.config.logger.info("not adding a schema")
return True
alembic_cfg = Config(self.config.alembic_config)
alembic_cfg.set_main_option('sqlalchemy.url', normal_user_pg_url)
self.config.logger.info('database schema section with %s',
normal_user_pg_url)
with PostgreSQLAlchemyManager(normal_user_pg_url,
self.config.logger,
autocommit=False) as db:
# Order matters below
db.turn_function_body_checks_off()
db.load_raw_sql('types')
db.load_raw_sql('procs')
# We need to commit to make a type visible for table creation
db.commit()
db.create_tables()
db.commit()
if not self.config.get('no_staticdata'):
self.import_staticdata(db)
db.commit()
command.stamp(alembic_cfg, "heads")
db.session.close()
# database owner section
self.config.logger.info('database extensions section with %s',
superuser_normaldb_pg_url)
with PostgreSQLAlchemyManager(superuser_normaldb_pg_url,
self.config.logger,
autocommit=False) as db:
db.set_table_owner(self.config.database_username)
db.set_default_owner(database_name, self.config.database_username)
db.set_grants(self.config) # config has user lists
return True
def main(self):
self.database_name = self.config['database_name']
if not self.database_name:
print "Syntax error: --database_name required"
return 1
self.no_schema = self.config.get('no_schema')
self.no_roles = self.config.get('no_roles')
self.force = self.config.get('force')
def connection_url():
url_template = 'postgresql://'
if self.database_username:
url_template += '%s' % self.database_username
if self.database_password:
url_template += ':%s' % self.database_password
url_template += '@'
if self.database_hostname:
url_template += '%s' % self.database_hostname
if self.database_port:
url_template += ':%s' % self.database_port
return url_template
self.database_username = self.config.get('database_superusername')
self.database_password = self.config.get('database_superuserpassword')
self.database_hostname = self.config.get('database_hostname')
self.database_port = self.config.get('database_port')
url_template = connection_url()
sa_url = url_template + '/%s' % 'postgres'
if self.config.unlogged:
@compiles(CreateTable)
def create_table(element, compiler, **kw):
text = compiler.visit_create_table(element, **kw)
text = re.sub("^\sCREATE(.*TABLE)",
lambda m: "CREATE UNLOGGED %s" %
m.group(1), text)
return text
with PostgreSQLAlchemyManager(sa_url, self.config.logger,
autocommit=False) as db:
if not db.min_ver_check(90200):
print 'ERROR - unrecognized PostgreSQL version: %s' % \
db.version_string()
print 'Only 9.2+ is supported at this time'
return 1
connection = db.engine.connect()
if self.config.get('dropdb'):
if 'test' not in self.database_name and not self.force:
confirm = raw_input(
'drop database %s [y/N]: ' % self.database_name)
if not confirm == "y":
logging.warn('NOT dropping table')
return 2
try:
# work around for autocommit behavior
connection.execute('commit')
connection.execute('DROP DATABASE %s' % self.database_name)
except exc.ProgrammingError, e:
if re.search(
'database "%s" does not exist' % self.database_name,
e.orig.pgerror.strip()):
# already done, no need to rerun
print "The DB %s doesn't exist" % self.database_name
try:
# work around for autocommit behavior
connection.execute('commit')
connection.execute("CREATE DATABASE %s ENCODING 'utf8'" %
self.database_name)
except exc.ProgrammingError, e:
if re.search(
'database "%s" already exists' % self.database_name,
e.orig.pgerror.strip()):
# already done, no need to rerun
print "The DB %s already exists" % self.database_name
return 0
raise
# already done, no need to rerun
print "The DB %s already exists" % self.database_name
return 0
raise
if self.no_roles:
print "Skipping role creation"
else:
db.create_roles(self.config)
connection.close()
# Reconnect to set up schema, types and procs
sa_url = url_template + '/%s' % self.database_name
alembic_cfg = Config(self.config.alembic_config)
alembic_cfg.set_main_option("sqlalchemy.url", sa_url)
with PostgreSQLAlchemyManager(sa_url, self.config.logger) as db:
connection = db.engine.connect()
db.setup_admin()
if self.no_schema:
db.commit()
return 0
# Order matters with what follows
db.create_types()
db.create_procs()
db.set_sequence_owner('breakpad_rw')
db.commit()
db.create_tables()
db.set_table_owner('breakpad_rw')
db.create_views()
