def test_reset_autoincrements(base_db):
engines = base_db.engines
shard_schemas = get_shard_schemas()
# A correctly set auto_increment.
key = 0
init_db(engines[key], key)
reset_tables = reset_invalid_autoincrements(engines[key],
shard_schemas[key], key, False)
assert len(reset_tables) == 0
# Ensure dry_run mode does not reset tables
key = 1
init_db(engines[key], key + 1)
reset_tables = reset_invalid_autoincrements(engines[key],
shard_schemas[key], key, True)
assert len(reset_tables) > 0
with pytest.raises(AssertionError):
verify_db(engines[key], shard_schemas[key], key)
reset_tables = reset_invalid_autoincrements(engines[key],
shard_schemas[key], key, False)
assert len(reset_tables) > 0
verify_db(engines[key], shard_schemas[key], key)
def test_reset_autoincrements(base_db):
engines = base_db.engines
shard_schemas = get_shard_schemas()
# A correctly set auto_increment.
key = 0
init_db(engines[key], key)
reset_tables = reset_invalid_autoincrements(engines[key],
shard_schemas[key], key,
False)
assert len(reset_tables) == 0
# Ensure dry_run mode does not reset tables
key = 1
init_db(engines[key], key + 1)
reset_tables = reset_invalid_autoincrements(engines[key],
shard_schemas[key], key,
True)
assert len(reset_tables) > 0
with pytest.raises(AssertionError):
verify_db(engines[key], shard_schemas[key], key)
reset_tables = reset_invalid_autoincrements(engines[key],
shard_schemas[key], key,
False)
assert len(reset_tables) > 0
verify_db(engines[key], shard_schemas[key], key)
def test_verify_db(base_db):
engines = base_db.engines
shard_schemas = get_shard_schemas()
# A correctly set auto_increment.
key = 0
init_db(engines[key], key)
verify_db(engines[key], shard_schemas[key], key)
# An incorrectly set auto_increment.
key = 1
init_db(engines[key], key + 1)
with pytest.raises(AssertionError):
verify_db(engines[key], shard_schemas[key], key)
def test_verify_db(base_db):
engines = base_db.engines
shard_schemas = get_shard_schemas()
# A correctly set auto_increment.
key = 0
init_db(engines[key], key)
verify_db(engines[key], shard_schemas[key], key)
# An incorrectly set auto_increment.
key = 1
init_db(engines[key], key + 1)
with pytest.raises(AssertionError):
verify_db(engines[key], shard_schemas[key], key)
def main():
maybe_enable_rollbar()
database_hosts = config.get_required("DATABASE_HOSTS")
database_users = config.get_required("DATABASE_USERS")
# Do not include disabled shards since application services do not use them.
engine_manager = EngineManager(database_hosts,
database_users,
include_disabled=False)
for host in database_hosts:
for shard in host["SHARDS"]:
if shard.get("DISABLED"):
continue
key = int(shard["ID"])
engine = engine_manager.engines[key]
schema = shard["SCHEMA_NAME"]
print("Verifying database: {}".format(schema))
verify_db(engine, schema, key)
def main(target_hostname, host_ip):
maybe_enable_rollbar()
database_hosts = config.get_required("DATABASE_HOSTS")
database_users = config.get_required("DATABASE_USERS")
engine_manager = EngineManager(
database_hosts, database_users, include_disabled=True
)
for host in database_hosts:
if target_hostname is not None and host["HOSTNAME"] != target_hostname:
continue
for shard in host["SHARDS"]:
key = shard["ID"]
assert isinstance(key, int)
hostname = host["HOSTNAME"]
connect_to = host_ip if host_ip else hostname
mysql_user = database_users[hostname]["USER"]
mysql_password = database_users[hostname]["PASSWORD"]
base_uri = build_uri(
username=mysql_user,
password=mysql_password,
hostname=connect_to,
port=host["PORT"],
database_name="",
)
base_engine = sqlalchemy.create_engine(
base_uri,
poolclass=ForceStrictModePool,
connect_args={"binary_prefix": True},
)
schema_name = shard["SCHEMA_NAME"]
print("Setting up database: {}".format(schema_name))
# Create the database IF needed.
base_engine.execute(
"CREATE DATABASE IF NOT EXISTS {} DEFAULT CHARACTER "
"SET utf8mb4 DEFAULT COLLATE utf8mb4_general_ci;".format(schema_name)
)
engine = engine_manager.engines[int(key)]
if engine.has_table("alembic_version"):
# Verify alembic revision
(current_revision,) = engine.execute(
"SELECT version_num from alembic_version"
).fetchone()
assert (
current_revision
), "Need current revision in alembic_version table."
print(
"Already revisioned by alembic version: {}".format(current_revision)
)
else:
# Initialize shards, stamp alembic revision
print("Initializing database.")
init_db(engine, int(key))
alembic_ini_filename = os.environ.get("ALEMBIC_INI_PATH", "alembic.ini")
assert os.path.isfile(
alembic_ini_filename
), "Must have alembic.ini file at {}".format(alembic_ini_filename)
alembic_cfg = alembic.config.Config(alembic_ini_filename)
# Alembic option values need to be strings.
alembic_cfg.set_main_option("shard_id", str(key))
print("Stamping with alembic revision.")
alembic.command.stamp(alembic_cfg, "head")
# Verify the database has been set up with correct auto_increments.
print("Verifying database.")
verify_db(engine, schema_name, int(key))
print("Finished setting up database.\n")
