Add MySQL support

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 (<user>:<password>@<host>:<port>/<database>)",
     )
     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()