Update the tool update_song.py

tools/update_song.py

@@ -1,11 +1,17 @@
+import os
 import sqlite3
 
+
+SONG_DATABASE_PATH = './arcsong.db'
+SERVER_DATABASE_PATH = './arcaea_database.db'
+
+
 class Connect():
     # 数据库连接类，上下文管理
 
-    def __init__(self, file_path='./arcaea_database.db'):
+    def __init__(self, file_path=SERVER_DATABASE_PATH):
         """
-            数据库连接，默认连接arcaea_database.db\ 
+            数据库连接\ 
             接受：文件路径\ 
             返回：sqlite3连接操作对象
         """
@@ -28,33 +34,81 @@ class Connect():
         return True
 
 
-def insert(cursor, song_id, name, a, b, c, d):
+def insert(cursor, song_id, name, a, b, c, d, update_type=0):
     '''Insert a new song into database.'''
-    cursor.execute(
-        '''select exists(select * from chart where song_id=?)''', (song_id, ))
-    if cursor.fetchone()[0]:
-        return None
+    if update_type == 0 or update_type == 1:
+        cursor.execute(
+            '''select exists(select * from chart where song_id=?)''', (song_id, ))
+        if cursor.fetchone()[0]:
+            if update_type == 0:
+                # 重复则不更新，以服务端数据库数据为准
+                return
+            elif update_type == 1:
+                # 重复则更新，以`arcsong.db`数据为准
+                cursor.execute('''update chart set name=?, rating_pst=?, rating_prs=?, rating_ftr=?, rating_byn=? where song_id=?''',
+                               (name, a, b, c, d, song_id))
+                return
+
     cursor.execute(
         '''insert into chart values (?,?,?,?,?,?)''', (song_id, name, a, b, c, d))
 
 
-def old_to_new():
-    '''Update old database to new database.'''
-    with Connect('./arcsong.db') as c:
-        c.execute(
-            '''select sid, name_en, rating_pst, rating_prs, rating_ftr, rating_byn from songs''')
-        data = c.fetchall()
+def from_song_datebase():
+    '''Get song data from song database and insert them into server's database.'''
+    with Connect(SONG_DATABASE_PATH) as c:
+        c.execute('''select name from sqlite_master where type="table"''')
+        tables = [x[0] for x in c.fetchall()]
+        if 'songs' in tables:
+            c.execute(
+                '''select sid, name_en, rating_pst, rating_prs, rating_ftr, rating_byn from songs''')
+            data = c.fetchall()
+        elif 'charts' in tables:
+            c.execute(
+                '''select song_id, rating_class, name_en, rating from charts''')
+            songs = {}
+            for song_id, rating_class, name_en, rating in c.fetchall():
+                if song_id not in songs:
+                    songs[song_id] = [-1, -1, -1, -1, name_en]
+                songs[song_id][rating_class] = rating
+
+            data = [(x, y[-1], y[0], y[1], y[2], y[3])
+                    for x, y in songs.items()]
+        else:
+            print('Error: Cannot find table `songs` or `charts` in the database.')
+            return
+
+    # 用户确认更新方式
+    update_type = 0
+    x = input('Type a number to decide the update type:\n0: Do not update if the song already exists in the server database.\n1: Update even if the song already exists in the server database.\n2: Clear chart data in the server database and then update.\nYour choice: ')
+    x = x.strip()
+    if x not in ('0', '1', '2'):
+        print('Error: Invalid input.')
+        return
+    update_type = int(x)
+
     with Connect() as c:
+        if update_type == 2:
+            # 清空数据表后更新
+            c.execute('''delete from chart''')
         for x in data:
-            insert(c, x[0], x[1], x[2], x[3], x[4], x[5])
+            insert(c, x[0], x[1], x[2], x[3], x[4], x[5], update_type)
+
+    print('Seems to be done.')
+
+
+def check_file():
+    if not os.path.isfile(SONG_DATABASE_PATH) or not os.path.isfile(SERVER_DATABASE_PATH):
+        print('Error: Files cannot be found.')
+        print('Note: Please make sure that both `arcsong.db` and `arcaea_server.db` are in this directory.')
+        return False
+    return True
 
 
 def main():
-    old_to_new()
+    if check_file():
+        from_song_datebase()
 
 
 if __name__ == '__main__':
     main()
-    print('Done.')
-    input()
-    exit()
+    input('Press `Enter` key to exit.')