def load_svt_extra_db(
engine: Engine,
svtExtras: list[MstSvtExtra]) -> None: # pragma: no cover
svtExtra_db_data = [svtExtra.dict() for svtExtra in svtExtras]
with engine.begin() as conn:
recreate_table(conn, mstSvtExtra)
conn.execute(mstSvtExtra.insert(), svtExtra_db_data)
def migrate_db(engine: Engine) -> None:
config_path = Path(__file__).resolve().parent.parent / "alembic.ini"
alembic_cfg = Config(config_path)
with engine.begin() as connection:
# check if database has not been migrated yet with alembic
if has_tables(connection) and not has_revision(connection):
# stamp database with the initial revision id, representing pre-alembic database schema
command.stamp(alembic_cfg, "c54a916ec6c8")
command.upgrade(alembic_cfg, "head")
def load_skill_td_lv(engine: Engine,
master_folder: DirectoryPath) -> None: # pragma: no cover
with open(master_folder / "mstBuff.json", "rb") as fp:
mstBuff = {buff["id"]: buff for buff in orjson.loads(fp.read())}
with open(master_folder / "mstFunc.json", "rb") as fp:
mstFunc = {func["id"]: func for func in orjson.loads(fp.read())}
mstFuncGroup = defaultdict(list)
with open(master_folder / "mstFuncGroup.json", "rb") as fp:
for funcGroup in orjson.loads(fp.read()):
mstFuncGroup[funcGroup["funcId"]].append(funcGroup)
with open(master_folder / "mstSkillLv.json", "rb") as fp:
mstSkillLv_data = orjson.loads(fp.read())
with open(master_folder / "mstTreasureDeviceLv.json", "rb") as fp:
mstTreasureDeviceLv_data = orjson.loads(fp.read())
def get_func_entity(func_id: int) -> dict[Any, Any]:
func_entity = {
"mstFunc": mstFunc[func_id],
"mstFuncGroup": mstFuncGroup.get(func_id, []),
}
if (func_entity["mstFunc"]["funcType"] not in FUNC_VALS_NOT_BUFF
and func_entity["mstFunc"]["vals"]
and func_entity["mstFunc"]["vals"][0] in mstBuff):
func_entity["mstFunc"]["expandedVals"] = [{
"mstBuff":
mstBuff[func_entity["mstFunc"]["vals"][0]]
}]
else:
func_entity["mstFunc"]["expandedVals"] = []
return func_entity
for skillLv in mstSkillLv_data:
skillLv["expandedFuncId"] = [
get_func_entity(func_id) for func_id in skillLv["funcId"]
if func_id in mstFunc
]
for treasureDeviceLv in mstTreasureDeviceLv_data:
treasureDeviceLv["expandedFuncId"] = [
get_func_entity(func_id) for func_id in treasureDeviceLv["funcId"]
if func_id in mstFunc
]
with engine.begin() as conn:
recreate_table(conn, mstSkillLv)
conn.execute(mstSkillLv.insert(), mstSkillLv_data)
recreate_table(conn, mstTreasureDeviceLv)
conn.execute(mstTreasureDeviceLv.insert(), mstTreasureDeviceLv_data)
def run_alembic_command(engine: Engine, command: str, command_kwargs: Dict[str, Any]) -> str:
command_func = ALEMBIC_COMMAND_MAP[command]
stdout = StringIO()
alembic_cfg = build_alembic_config(engine)
with engine.begin() as connection:
alembic_cfg.attributes["connection"] = connection
with contextlib.redirect_stdout(stdout):
command_func(alembic_cfg, **command_kwargs)
return stdout.getvalue()
def change_functions_ownership(engine: Engine, database: str,
target_role: str):
stmt = f"SHOW USER FUNCTIONS IN DATABASE {database}"
with engine.begin() as tx:
for object in [dict(x) for x in tx.execute(stmt).fetchall()]:
func_handle = object["arguments"].split(" RETURN ")[0]
schema = object["schema_name"]
tx.execute(
f"GRANT OWNERSHIP ON FUNCTION {database}.{schema}.{func_handle} "
f"TO ROLE {target_role} REVOKE CURRENT GRANTS").fetchall()
def load_script_list(engine: Engine,
repo_folder: DirectoryPath) -> None: # pragma: no cover
script_list_file = (repo_folder / "ScriptActionEncrypt" /
ScriptFileList.name / f"{ScriptFileList.name}.txt")
db_data: list[dict[str, Union[int, str, None]]] = []
if script_list_file.exists():
with open(script_list_file, encoding="utf-8") as fp:
script_list = [line.strip() for line in fp.readlines()]
with open(repo_folder / "master" / "mstQuest.json", "rb") as bfp:
mstQuest = orjson.loads(bfp.read())
questId = {quest["id"] for quest in mstQuest}
scriptQuestId = {
quest["scriptQuestId"]: quest["id"]
for quest in mstQuest if quest["scriptQuestId"] != 0
}
for script in script_list:
script_name = script.removesuffix(".txt")
quest_ids: list[Optional[int]] = []
phase: Optional[int] = None
sceneType: Optional[int] = None
if len(script) == 14 and script[0] in ("0", "9"):
script_int = int(script_name[:-2])
sceneType = int(script_name[-1])
phase = int(script_name[-2])
if script_int in scriptQuestId:
quest_ids.append(scriptQuestId[script_int])
if script_int in questId:
quest_ids.append(script_int)
if not quest_ids:
quest_ids.append(None)
for quest_id in quest_ids:
db_data.append({
"scriptFileName": script_name,
"questId": quest_id,
"phase": phase,
"sceneType": sceneType,
})
with engine.begin() as conn:
recreate_table(conn, ScriptFileList)
conn.execute(ScriptFileList.insert(), db_data)
def aggregate_report(bookings_table: MappedTable, report_table: MappedTable,
engine: Engine):
"""
Aggregate bookings table into a monthly report, and replace report table content with it.
Args:
bookings_table: source bookings table mapper
report_table: destination report table mapper
engine: sqlalchemy engine
Notes:
Make the operation transactional.
"""
with engine.begin() as connection:
# Truncate already-existing report table
_run_truncate_query(report_table, connection)
# Recreate report
_run_aggregate_report_query(bookings_table, report_table, connection)
def change_objects_ownership(engine: Engine, database: str,
target_role: str) -> None:
stmt = Select(
[
literal_column("table_type"),
literal_column("table_schema"),
literal_column("table_name"),
],
from_obj=text(f"{database}.INFORMATION_SCHEMA.TABLES"),
whereclause=literal_column("table_owner") == "DBT_PRODUCTION",
)
with engine.begin() as tx:
rp = tx.execute(stmt)
objects = [(
"TABLE" if object_type == "BASE TABLE" else object_type,
schema,
object_name,
) for object_type, schema, object_name in rp.fetchall()]
for object_type, schema, object_name in objects:
tx.execute(
f"GRANT OWNERSHIP ON {object_type} {database}.{schema}.{object_name} TO ROLE {target_role} REVOKE CURRENT GRANTS"
).fetchall()
def get_current_head(connectable: engine.Engine) -> set:
with connectable.begin() as connection:
context = migration.MigrationContext.configure(connection)
current_head = set(context.get_current_heads())
return current_head
def fetch_user_stickers(db_pool: engine.Engine) -> List[str]:
with db_pool.begin() as conn:
old_image_ids = conn.execute(user_image_query).fetchall()
return [str(row).split("'")[1] for row in old_image_ids]