Merge pull request 'Add MySQL support' (#2) from Floatin/AQUA-to-ARTEMiS:mysql-support into main

Reviewed-on: #2

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()