diff --git a/aqua_importer.py b/aqua_importer.py index 86d53f8..bcffa21 100644 --- a/aqua_importer.py +++ b/aqua_importer.py @@ -10,7 +10,6 @@ from sqlalchemy.exc import SQLAlchemyError from logging.handlers import TimedRotatingFileHandler from typing import Any, Dict, Optional import yaml -import yaml import argparse import logging import coloredlogs @@ -27,9 +26,12 @@ from titles.ongeki.const import OngekiConstants class AquaData: def __init__(self, aqua_db_path: str) -> None: - self.__url = f"sqlite:///{aqua_db_path}" + if '@' in aqua_db_path: + self.__url = f"mysql+pymysql://{aqua_db_path}" + else: + self.__url = f"sqlite:///{aqua_db_path}" - self.__engine = create_engine(self.__url, pool_recycle=3600) + self.__engine = create_engine(self.__url, pool_recycle=3600, echo=False) # self.inspector = reflection.Inspector.from_engine(self.__engine) session = sessionmaker(bind=self.__engine) @@ -114,17 +116,47 @@ 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 not os.path.isfile(f'{aqua_folder}/application.properties'): + self.logger.error("Could not locate AQUA application.properties file!") + exit(1) + + with open(f'{aqua_folder}/application.properties') as file: + lines = file.readlines() + + properties = {} + for line in lines: + line = line.strip() + if not line or line.startswith('#'): + continue + parts = line.split('=') + if len(parts) >= 2: + key = parts[0].strip() + value = '='.join(parts[1:]).strip() + properties[key] = value + + db_driver = properties.get('spring.datasource.driver-class-name') + if 'sqlite' in db_driver: + aqua_db_path = None + db_url = properties.get('spring.datasource.url').split('sqlite:')[1] + temp = os.path.join(f'{aqua_folder}/{db_url}') 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) + + elif 'mysql' in db_driver or 'mariadb' in db_driver: + self.use_mysql = True + db_username = properties.get('spring.datasource.username') + db_password = properties.get('spring.datasource.password') + db_url = properties.get('spring.datasource.url').split('?')[0].split('//')[1] + self.aqua = AquaData(f'{db_username}:{db_password}@{db_url}') + + else: + self.logger.error("Unknown database type!") def get_user_id(self, luid: str): user_id = self.data.card.get_user_id_from_card(access_code=luid) @@ -164,14 +196,15 @@ 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 not self.use_mysql: + 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(): @@ -706,7 +739,7 @@ def main(): parser.add_argument( "aqua_folder_path", type=str, - help="Absolute folder path to AQUA /data folder, where db.sqlite is located in", + help="The absolute folder path to the folder where AQUA is located, where the data folder and the application.properties file should be located.", ) args = parser.parse_args() @@ -718,6 +751,5 @@ def main(): importer.import_chuni() importer.import_ongeki() - if __name__ == "__main__": - main() + main() \ No newline at end of file