def setup_test_database():
"""
Update the test DB to the latest version. The migration code is able to
handle both updating an existing, pre-migration code DB (local tests)
and creating a new DB (automated tests)
"""
alembic_main(["--raiseerr", "upgrade", "head"])
def create_db(uri=None):
"""Create a database specified by uri and setup tables.
:arg str uri: either None or a valid uri; if None, uses ``SQLALCHEMY_URL``
environment variable
:raises sqlalchemy.exc.ProgrammingError: if database already exists
:raises SqlAlchemyUrlNotSpecified: if ``SQLALCHEMY_URL`` has no value
Note: This is mysql-specific.
"""
if uri is None:
uri = get_sqlalchemy_url()
sa_url = make_url(uri)
db_to_create = sa_url.database
if hasattr(sa_url, "set"):
# SQLAlchemy 1.4: immutable URL object
sa_url_no_db = sa_url.set(database="")
else:
# SQLAlchemy 1.3: mutate in place
sa_url_no_db = sa_url
sa_url_no_db.database = None
engine = create_engine(sa_url_no_db)
with engine.connect() as conn:
conn.execute(f"CREATE DATABASE {db_to_create} CHARACTER SET = 'utf8'")
alembic_main(["stamp", "base"])
alembic_main(["upgrade", "head"])
async def test_42cbae986650_upgrade(resource_path):
# before "Add policy_id and revoke columns to request table" migration
alembic_main(["--raiseerr", "downgrade", "c0a92da5ac69"])
# insert a request
uuid = "571c6a1a-f21f-11ea-adc1-0242ac120002"
date = str(datetime.now())
sql_resource_path = resource_path.replace("'", "''") # escape single quotes
insert_stmt = f"INSERT INTO requests(\"request_id\", \"username\", \"resource_path\", \"resource_id\", \"resource_display_name\", \"status\", \"created_time\", \"updated_time\") VALUES ('{uuid}', 'username', '{sql_resource_path}', 'my_resource', 'My Resource', 'DRAFT', '{date}', '{date}')"
await db.scalar(db.text(insert_stmt))
# check that the request data was inserted correctly
data = await db.all(db.text("SELECT * FROM requests"))
request = {k: str(v) for row in data for k, v in row.items()}
assert request == {
"request_id": uuid,
"username": "******",
"resource_id": "my_resource",
"resource_display_name": "My Resource",
"status": "DRAFT",
"resource_path": resource_path,
"created_time": date,
"updated_time": date,
}
# run "Add policy_id and revoke columns to request table" migration
alembic_main(["--raiseerr", "upgrade", "42cbae986650"])
# check that the migration updated the request data correctly
data = await db.all(db.text("SELECT * FROM requests"))
assert len(data) == 1
request = {k: v for k, v in data[0].items()}
assert resource_path not in request
assert request["policy_id"] == get_auto_policy_id_for_resource_path(resource_path)
assert request["revoke"] == False
def setup_tables(engine):
# Stamp the latest alembic version
alembic_main(["stamp", "base"])
alembic_main(["upgrade", "head"])
# Setup table data
_setup_table_data(engine)
def __call__(self, args):
current_directory = os.curdir
try:
os.chdir(args.application.root_path)
alembic_ini = join(args.application.root_path, 'alembic.ini')
alembic_main(argv=['--config', alembic_ini] + args.alembic_args)
finally:
os.chdir(current_directory)
def launch(self):
current_directory = os.curdir
try:
os.chdir(self.args.application.root_path)
alembic_ini = join(self.args.application.root_path, 'alembic.ini')
alembic_main(argv=['--config', alembic_ini] +
self.args.alembic_args)
finally:
os.chdir(current_directory)
def setup_test_database():
"""
At teardown, restore original config and reset test DB.
"""
saved_config = copy.deepcopy(config._configs)
alembic_main(["--raiseerr", "upgrade", "head"])
yield
# restore old configs
config.update(saved_config)
if not config["TEST_KEEP_DB"]:
alembic_main(["--raiseerr", "downgrade", "base"])
def clean_db():
"""
Before each test, delete all existing requests from the DB
"""
# The code below doesn't work because of this issue
# https://github.com/encode/starlette/issues/440, so for now reset
# using alembic.
# pytest-asyncio = "^0.14.0"
# from requestor.models import Request as RequestModel
# @pytest.mark.asyncio
# async def clean_db():
# await RequestModel.delete.gino.all()
# yield
alembic_main(["--raiseerr", "downgrade", "base"])
alembic_main(["--raiseerr", "upgrade", "head"])
yield
def main(argv=None, **kwargs):
if not argv:
argv = sys.argv
# clean up args
argv.pop(0)
if len(argv) >= 1 and argv[0] == 'initdb':
print "Initializing db."
initialize_db()
print "DB initialized."
exit(1)
# default config arg, if not already set
if not '-c' in argv:
argv.insert(0, DEFAULT_ALEMBIC_INI_PATH)
argv.insert(0, '-c')
alembic_main(argv, **kwargs)
def main(argv=None, **kwargs):
if not argv:
argv = sys.argv
# clean up args
argv.pop(0)
if len(argv) >= 1 and argv[0] == 'initdb':
print "Initializing db."
initialize_db(PITHOS_BACKEND_DB_CONNECTION)
print "DB initialized."
exit(1)
# default config arg, if not already set
if not '-c' in argv:
argv.insert(0, DEFAULT_ALEMBIC_INI_PATH)
argv.insert(0, '-c')
alembic_main(argv, **kwargs)
def run_alembic(args: List[str]) -> None:
"""
Forward arguments to alembic
"""
from alembic.config import main as alembic_main
dbkey = "ERT_STORAGE_DATABASE_URL"
dburl = os.getenv(dbkey)
if dburl is None:
sys.exit(
f"Environment variable '{dbkey}' not set.\n"
"It needs to point to a PostgreSQL server for alembic to work.")
if not dburl.startswith("postgresql"):
sys.exit(
f"Environment variable '{dbkey}' does not point to a postgresql database.\n"
"Only PostgreSQL is supported for alembic migrations at the moment.\n"
f"Its value is: {dburl}")
argv = [
"-c",
os.path.join(os.path.dirname(__file__), "_alembic", "alembic.ini"),
*args,
]
try:
alembic_main(argv=argv, prog="ert-storage alembic")
except FileNotFoundError as exc:
if os.path.basename(exc.filename) == "script.py.mako":
sys.exit(
f"\nAlembic could not find 'script.py.mako' in location:\n"
f"\n{exc.filename}\n\n"
"This is most likely because you've installed ert-storage without --edit mode\n"
"Reinstall ert-storage with: pip install -e <path to ert-storage>"
)
else:
raise
sys.exit(0)
def create_db(uri=None):
"""Create a database specified by uri and setup tables.
:arg str uri: either None or a valid uri; if None, uses ``SQLALCHEMY_URL``
environment variable
:raises sqlalchemy.exc.ProgrammingError: if database already exists
:raises SqlAlchemyUrlNotSpecified: if ``SQLALCHEMY_URL`` has no value
Note: This is mysql-specific.
"""
if uri is None:
uri = get_sqlalchemy_url()
sa_url = make_url(uri)
db_to_create = sa_url.database
sa_url.database = None
engine = create_engine(sa_url)
engine.execute(
"CREATE DATABASE {} CHARACTER SET = 'utf8'".format(db_to_create))
alembic_main(["stamp", "base"])
alembic_main(["upgrade", "head"])
def migrations(_args):
migrations_path = pkg_resources.resource_filename("antalla", "migrations")
os.chdir(migrations_path)
alembic_main(sys.argv[2:], prog="antalla migrations")
def alembic(*args):
if not db_exists():
raise exceptions.DBError("Database do not exists")
aargs = ["--config", conf.settings.MIGRATIONS_SETTINGS] + list(args)
return alembic_main(aargs, "corral")
def launch(self):
alembic_ini = join(self.args.application.root_path, 'alembic.ini')
alembic_main(argv=['--config', alembic_ini] + self.args.alembic_args)
def alembic(*args):
if not db_exists():
raise exceptions.DBError("Database do not exists")
aargs = ["--config", conf.settings.MIGRATIONS_SETTINGS] + list(args)
return alembic_main(aargs, "corral")
def _bootstrap():
"""Create all of the models based on the declaritive base for this service."""
alembic_config = os.path.join(os.path.split(__file__)[0], '..', 'alembic.ini')
alembic_main(['-c', alembic_config, 'upgrade', 'head'])
def execute(self, args):
from alembic.config import main as alembic_main
config_path = get_config('alembic.alembic_ini')
argv = ['-c', config_path] + args
alembic_main(prog='koschei-admin alembic', argv=argv)
def execute(self, args):
from alembic.config import main as alembic_main
config_path = get_config('alembic.alembic_ini')
argv = ['-c', config_path] + args
alembic_main(prog='koschei-admin alembic', argv=argv)