import time
from ... import exc
from ... import text
from ...testing.provision import create_db
from ...testing.provision import drop_db
from ...testing.provision import log
from ...testing.provision import temp_table_keyword_args
@create_db.for_db("postgresql")
def _pg_create_db(cfg, eng, ident):
template_db = cfg.options.postgresql_templatedb
with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
try:
_pg_drop_db(cfg, conn, ident)
except Exception:
pass
if not template_db:
template_db = conn.scalar("select current_database()")
attempt = 0
while True:
try:
conn.execute(
"CREATE DATABASE %s TEMPLATE %s" % (ident, template_db)
)
except exc.OperationalError as err:
attempt += 1
if attempt >= 3:
raise
if "accessed by other users" in str(err):
log.info(
"Waiting to create %s, URI %r, "
"template DB %s is in use sleeping for .5",
ident,
eng.url,
template_db,
)
time.sleep(0.5)
except:
raise
else:
break
@drop_db.for_db("postgresql")
def _pg_drop_db(cfg, eng, ident):
with eng.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
conn.execute(
text(
"select pg_terminate_backend(pid) from pg_stat_activity "
"where usename=current_user and pid != pg_backend_pid() "
"and datname=:dname"
),
dname=ident,
)
conn.execute("DROP DATABASE %s" % ident)
@temp_table_keyword_args.for_db("postgresql")
def _postgresql_temp_table_keyword_args(cfg, eng):
return {"prefixes": ["TEMPORARY"]}
|