From 7af88d7baa791741d4d4eb9e8473b5eb7fcaab34 Mon Sep 17 00:00:00 2001 From: Floatin Date: Fri, 20 Oct 2023 07:26:22 +0000 Subject: [PATCH] Add MySQL support --- aqua_importer.py | 59 +++++++++++++++++++++++++++++------------------- 1 file changed, 36 insertions(+), 23 deletions(-) diff --git a/aqua_importer.py b/aqua_importer.py index 86d53f8..551313d 100644 --- a/aqua_importer.py +++ b/aqua_importer.py @@ -27,10 +27,13 @@ from titles.ongeki.const import OngekiConstants class AquaData: def __init__(self, aqua_db_path: str) -> None: - self.__url = f"sqlite:///{aqua_db_path}" + if use_mysql: + self.__engine = create_engine("mysql+pymysql://" + aqua_db_path, echo=False) + else: + self.__url = f"sqlite:///{aqua_db_path}" - self.__engine = create_engine(self.__url, pool_recycle=3600) - # self.inspector = reflection.Inspector.from_engine(self.__engine) + self.__engine = create_engine(self.__url, pool_recycle=3600) + # self.inspector = reflection.Inspector.from_engine(self.__engine) session = sessionmaker(bind=self.__engine) self.inspect = inspect(self.__engine) @@ -114,17 +117,20 @@ class Importer: coloredlogs.install(level="INFO", logger=self.logger, fmt=log_fmt_str) self.logger.initialized = True - aqua_db_path = None - if os.path.exists(aqua_folder): - temp = os.path.join(aqua_folder, "db.sqlite") - if os.path.isfile(temp): - aqua_db_path = temp + if use_mysql: + self.aqua = AquaData(aqua_folder) + else: + aqua_db_path = None + if os.path.exists(aqua_folder): + temp = os.path.join(aqua_folder, "db.sqlite") + if os.path.isfile(temp): + aqua_db_path = temp - if not aqua_db_path: - self.logger.error("Could not locate AQUA db.sqlite file!") - exit(1) + if not aqua_db_path: + self.logger.error("Could not locate AQUA db.sqlite file!") + exit(1) - self.aqua = AquaData(aqua_db_path) + self.aqua = AquaData(aqua_db_path) def get_user_id(self, luid: str): user_id = self.data.card.get_user_id_from_card(access_code=luid) @@ -164,14 +170,16 @@ class Importer: card_id: int, ) -> Dict: row = row._asdict() - for column in datetime_columns: - ts = row[column["name"]] - if ts is None: - continue - # actuall remove the last 3 zeros for the correct timestamp - fixed_ts = int(str(ts)[:-3]) - # save the datetim object in the dict - row[column["name"]] = datetime.fromtimestamp(fixed_ts) + + if use_mysql is False: + for column in datetime_columns: + ts = row[column["name"]] + if ts is None: + continue + # actuall remove the last 3 zeros for the correct timestamp + fixed_ts = int(str(ts)[:-3]) + # save the datetim object in the dict + row[column["name"]] = datetime.fromtimestamp(fixed_ts) tmp = {} for k, v in row.items(): @@ -704,16 +712,21 @@ def main(): "--config", "-c", type=str, help="Config directory to use", default="config" ) parser.add_argument( - "aqua_folder_path", + "aqua_data_path", type=str, - help="Absolute folder path to AQUA /data folder, where db.sqlite is located in", + help="Absolute folder path to AQUA /data folder, where db.sqlite is located in. You can also enter the Mysql database connection address (:@:/)", ) args = parser.parse_args() core_cfg = CoreConfig() core_cfg.update(yaml.safe_load(open(f"{args.config}/core.yaml"))) - importer = Importer(core_cfg, args.config, args.aqua_folder_path) + global use_mysql + use_mysql = False + if '@' in args.aqua_data_path: + use_mysql = True + + importer = Importer(core_cfg, args.config, args.aqua_data_path) importer.import_chuni() importer.import_ongeki()