mai2: add opts to reader

titles/mai2/read.py

@@ -1,20 +1,16 @@
-from decimal import Decimal
-import logging
 import os
 import re
 import xml.etree.ElementTree as ET
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 from Crypto.Cipher import AES
 import zlib
 import codecs
 
 from core.config import CoreConfig
-from core.data import Data
 from read import BaseReader
 from titles.mai2.const import Mai2Constants
 from titles.mai2.database import Mai2Data
 
-
 class Mai2Reader(BaseReader):
     def __init__(
         self,
@@ -46,10 +42,11 @@ class Mai2Reader(BaseReader):
 
             for dir in data_dirs:
                 self.logger.info(f"Read from {dir}")
-                await self.get_events(f"{dir}/event")
+                this_opt_id = await self.read_opt_info(dir)
+                await self.get_events(f"{dir}/event", this_opt_id)
                 await self.disable_events(f"{dir}/information", f"{dir}/scoreRanking")
-                await self.read_music(f"{dir}/music")
+                await self.read_music(f"{dir}/music", this_opt_id)
-                await self.read_tickets(f"{dir}/ticket")
+                await self.read_tickets(f"{dir}/ticket", this_opt_id)
         
         else:
             if not os.path.exists(f"{self.bin_dir}/tables"):
@@ -179,7 +176,7 @@ class Mai2Reader(BaseReader):
             self.logger.warning("Failed load table content, skipping")
             return
 
-    async def get_events(self, base_dir: str) -> None:
+    async def get_events(self, base_dir: str, opt_id: int = None) -> None:
         self.logger.info(f"Reading events from {base_dir}...")
 
         for root, dirs, files in os.walk(base_dir):
@@ -193,7 +190,7 @@ class Mai2Reader(BaseReader):
                         event_type = int(troot.find("infoType").text)
 
                         await self.data.static.put_game_event(
-                            self.version, event_type, id, name
+                            self.version, event_type, id, name, opt_id
                         )
                         self.logger.info(f"Added event {id}...")
 
@@ -255,7 +252,7 @@ class Mai2Reader(BaseReader):
             await self.data.static.toggle_game_event(self.version, event_id, toggle=False)
             self.logger.info(f"Disabled event {event_id}...")
 
-    async def read_music(self, base_dir: str) -> None:
+    async def read_music(self, base_dir: str, opt_id: int = None) -> None:
         self.logger.info(f"Reading music from {base_dir}...")
 
         for root, dirs, files in os.walk(base_dir):
@@ -296,13 +293,14 @@ class Mai2Reader(BaseReader):
                                         added_ver,
                                         diff_num,
                                         note_designer,
+                                        opt_id
                                     )
 
                                     self.logger.info(
                                         f"Added music id {song_id} chart {chart_id}"
                                     )
 
-    async def read_tickets(self, base_dir: str) -> None:
+    async def read_tickets(self, base_dir: str, opt_id: int = None) -> None:
         self.logger.info(f"Reading tickets from {base_dir}...")
 
         for root, dirs, files in os.walk(base_dir):
@@ -317,7 +315,7 @@ class Mai2Reader(BaseReader):
                         price = int(troot.find("creditNum").text)
 
                         await self.data.static.put_game_ticket(
-                            self.version, id, ticket_type, price, name
+                            self.version, id, ticket_type, price, name, opt_id
                         )
                         self.logger.info(f"Added ticket {id}...")
 
@@ -341,3 +339,51 @@ class Mai2Reader(BaseReader):
         if scores is None or text is None:
             return
         # TODO
+
+    async def read_opt_info(self, directory: str) -> Optional[int]:
+        datacfg_file = os.path.join(directory, "DataConfig.xml")
+        if not os.path.exists(datacfg_file):
+            self.logger.warning(f"{datacfg_file} does not contain DataConfig.xml, opt info will not be read")
+            return None
+        
+        with open(datacfg_file, encoding="utf-8") as f:
+            troot = ET.fromstring(f.read())
+            
+            if troot.find("DataConfig/version") is None:
+                self.logger.warning(f"{directory}/DataConfig.xml contains no Version section, opt info will not be read")
+                return None
+            
+            ver_maj = troot.find("DataConfig/version/major")
+            ver_min = troot.find("DataConfig/version/minor")
+            ver_rel = troot.find("DataConfig/version/release")
+            cm_maj = troot.find("DataConfig/cardMakerVersion/major")
+            cm_min = troot.find("DataConfig/cardMakerVersion/minor")
+            cm_rel = troot.find("DataConfig/cardMakerVersion/release")
+            
+            if ver_maj is None: # Probably not worth checking that the other sections exist
+                self.logger.warning(f"{datacfg_file} contains no major item in the Version section, opt info will not be read")
+                return None
+
+            if ver_min is None: # Probably not worth checking that the other sections exist
+                self.logger.warning(f"{datacfg_file} contains no minor item in the Version section, opt info will not be read")
+                return None
+
+            if ver_rel is None: # Probably not worth checking that the other sections exist
+                self.logger.warning(f"{datacfg_file} contains no release item in the Version section, opt info will not be read")
+                return None
+            
+            opt_folder = os.path.basename(os.path.normpath(directory))
+            opt_id = await self.data.static.get_opt_by_version_folder(self.version, opt_folder)
+            
+            if not opt_id:
+                opt_id = await self.data.static.put_opt(self.version, opt_folder, int(ver_rel.text), int(cm_rel.text) if cm_rel else None)
+                if not opt_id:
+                    self.logger.error(f"Failed to put opt folder info for {opt_folder}")
+                    return None
+            else:
+                opt_id = opt_id['id']
+
+            self.logger.info(
+                f"Opt folder {opt_folder} (Database ID {opt_id}) contains v{ver_maj.text}.{ver_min.text}.{ver_rel.text} (cm v{cm_maj.text if cm_maj else 'None'}.{cm_min.text if cm_min else 'None'}.{cm_rel.text if cm_rel else 'None'})"
+            )
+            return opt_id

titles/mai2/schema/static.py

@@ -7,6 +7,7 @@ from sqlalchemy.schema import ForeignKey
 from sqlalchemy.sql import func, select
 from sqlalchemy.engine import Row
 from sqlalchemy.dialects.mysql import insert
+from sqlalchemy.sql.functions import coalesce
 from datetime import datetime
 
 opts = Table(
@@ -92,16 +93,17 @@ cards = Table(
 
 class Mai2StaticData(BaseData):
     async def put_game_event(
-        self, version: int, type: int, event_id: int, name: str
+        self, version: int, type: int, event_id: int, name: str, opt_id: int = None
     ) -> Optional[int]:
         sql = insert(event).values(
             version=version,
             type=type,
             eventId=event_id,
             name=name,
+            opt=coalesce(event.c.opt, opt_id)
         )
 
-        conflict = sql.on_duplicate_key_update(eventId=event_id)
+        conflict = sql.on_duplicate_key_update(eventId=event_id, opt=coalesce(event.c.opt, opt_id))
 
         result = await self.execute(conflict)
         if result is None:
@@ -154,6 +156,7 @@ class Mai2StaticData(BaseData):
         added_version: str,
         difficulty: float,
         note_designer: str,
+        opt_id: int = None
     ) -> None:
         sql = insert(music).values(
             version=version,
@@ -166,6 +169,7 @@ class Mai2StaticData(BaseData):
             addedVersion=added_version,
             difficulty=difficulty,
             noteDesigner=note_designer,
+            opt=coalesce(music.c.opt, opt_id)
         )
 
         conflict = sql.on_duplicate_key_update(
@@ -176,6 +180,7 @@ class Mai2StaticData(BaseData):
             addedVersion=added_version,
             difficulty=difficulty,
             noteDesigner=note_designer,
+            opt=coalesce(music.c.opt, opt_id)
         )
 
         result = await self.execute(conflict)
@@ -191,6 +196,7 @@ class Mai2StaticData(BaseData):
         ticket_type: int,
         ticket_price: int,
         name: str,
+        opt_id: int = None
     ) -> Optional[int]:
         sql = insert(ticket).values(
             version=version,
@@ -198,11 +204,10 @@ class Mai2StaticData(BaseData):
             kind=ticket_type,
             price=ticket_price,
             name=name,
+            opt=coalesce(ticket.c.opt, opt_id)
         )
 
-        conflict = sql.on_duplicate_key_update(price=ticket_price)
+        conflict = sql.on_duplicate_key_update(price=ticket_price, opt=coalesce(ticket.c.opt, opt_id))
-
-        conflict = sql.on_duplicate_key_update(price=ticket_price)
 
         result = await self.execute(conflict)
         if result is None:
@@ -247,12 +252,12 @@ class Mai2StaticData(BaseData):
             return None
         return result.fetchone()
 
-    async def put_card(self, version: int, card_id: int, card_name: str, **card_data) -> int:
+    async def put_card(self, version: int, card_id: int, card_name: str, opt_id: int = None, **card_data) -> int:
         sql = insert(cards).values(
-            version=version, cardId=card_id, cardName=card_name, **card_data
+            version=version, cardId=card_id, cardName=card_name, opt=coalesce(cards.c.opt, opt_id) **card_data
         )
 
-        conflict = sql.on_duplicate_key_update(**card_data)
+        conflict = sql.on_duplicate_key_update(opt=coalesce(cards.c.opt, opt_id), **card_data)
 
         result = await self.execute(conflict)
         if result is None:
@@ -282,3 +287,85 @@ class Mai2StaticData(BaseData):
         result = await self.execute(event.update(event.c.id == table_id).values(enabled=is_enable, startDate = start_date))
         if not result:
             self.logger.error(f"Failed to update event {table_id} - {is_enable} {start_date}")
+
+    async def put_opt(self, version: int, folder: str, sequence: int, cm_seq: int = None) -> Optional[int]:
+        sql = insert(opts).values(version=version, name=folder, sequence=sequence, cmReleaseVer=cm_seq)
+
+        conflict = sql.on_duplicate_key_update(sequence=sequence, whenRead=datetime.now())
+        
+        result = await self.execute(conflict)
+        if result is None:
+            self.logger.warning(f"Failed to insert opt! version {version} folder {folder} sequence {sequence}")
+            return None
+        return result.lastrowid
+
+    async def get_opt_by_version_folder(self, version: int, folder: str) -> Optional[Row]:        
+        result = await self.execute(opts.select(and_(
+            opts.c.version == version,
+            opts.c.name == folder,
+        )))
+        
+        if result is None:
+            return None
+        return result.fetchone()
+
+    async def get_opt_by_version_sequence(self, version: int, sequence: str) -> Optional[Row]:        
+        result = await self.execute(opts.select(and_(
+            opts.c.version == version,
+            opts.c.sequence == sequence,
+        )))
+        
+        if result is None:
+            return None
+        return result.fetchone()
+
+    async def get_opts_by_version(self, version: int) -> Optional[List[Row]]:        
+        result = await self.execute(opts.select(opts.c.version == version))
+        
+        if result is None:
+            return None
+        return result.fetchall()
+
+    async def get_opts_enabled_by_version(self, version: int) -> Optional[List[Row]]:        
+        result = await self.execute(opts.select(and_(
+            opts.c.version == version,
+            opts.c.isEnable == True,
+        )))
+        
+        if result is None:
+            return None
+        return result.fetchall()
+
+    async def get_latest_enabled_opt_by_version(self, version: int) -> Optional[Row]:        
+        result = await self.execute(
+            opts.select(and_(
+                opts.c.version == version,
+                opts.c.isEnable == True,
+            )).order_by(opts.c.sequence.desc())
+        )
+        
+        if result is None:
+            return None
+        return result.fetchone()
+
+    async def get_opts(self) -> Optional[List[Row]]:        
+        result = await self.execute(opts.select())
+        
+        if result is None:
+            return None
+        return result.fetchall()
+
+    async def get_opts(self) -> Optional[List[Row]]:        
+        result = await self.execute(opts.select())
+        
+        if result is None:
+            return None
+        return result.fetchall()
+
+    async def set_opt_enabled(self, opt_id: int, enabled: bool) -> bool:        
+        result = await self.execute(opts.update(opts.c.id == opt_id).values(isEnable=enabled))
+        
+        if result is None:
+            self.logger.error(f"Failed to set opt enabled status to {enabled} for opt {opt_id}")
+            return False
+        return True