65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
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"]}
|