[Fix] 修复数据库上传问题

1 files changed, 75 insertions, 30 deletions

diff --git a/ATRI/plugins/upload_setu.py b/ATRI/plugins/upload_setu.py
index 78158c3..c9dec60 100644
--- a/ATRI/plugins/upload_setu.py
+++ b/ATRI/plugins/upload_setu.py
@@ -1,29 +1,29 @@
 # -*- coding:utf-8 -*-
+import os
 import time
 import json
 import sqlite3
 import nonebot
 
+from pathlib import Path
 from nonebot import on_command, CommandSession
-from nonebot import session
 
 from ATRI.modules import response # type: ignore
 
 
 bot = nonebot.get_bot()
 master = bot.config.SUPERUSERS # type: ignore
-path_setu = f'ATRI\\splite\\setu\\'
-path_cloudmusic = f'ATRI\\splite\\cloudmusic\\'
 url = f'https://api.imjad.cn/pixiv/v1/?type=illust&id='
 
 
 @on_command('upload_setu', aliases = ['上传涩图'], only_to_me = False)
 async def _(session: CommandSession):
     if session.event.user_id in master:
+        start = time.perf_counter()
         msg = session.event.raw_message.split(' ', 2)
         print(msg)
         i_tpye = msg[1]
-        pid = msg[2]
+        pid = msg[2] 
 
         URL = url + pid
 
@@ -39,80 +39,125 @@ async def _(session: CommandSession):
         img = f'https://pixiv.cat/{pid}.jpg'
 
         if i_tpye == '正常':
-            con = sqlite3.connect(path_setu + 'normal.db')
+            if os.path.exists('ATRI/sqlite/setu/normal.db'):
+                print('数据文件存在！')
+            else:
+                await session.send('数据库不存在，将在3秒后开始构建...')
+                time.sleep(3)
+                await session.send('开始构建数据库！')
+                con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'normal.db')
+                cur = con.cursor()
+                cur.execute('CREATE TABLE normal(pid PID, title TITLE, tags TAGS, account ACCOUNT, name NAME, u_id UID, user_link USERLINK, img IMG, UNIQUE(pid, title, tags, account, name, u_id, user_link, img))')
+                con.commit()
+                cur.close()
+                con.close()
+                await session.send('完成')
+     
+            con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'normal.db')
             cur = con.cursor()
-            
             cur.execute('INSERT INTO normal VALUES ("%s", "%s", "%s", "%s", "%s", "%s", "%s", "%s")'%(pid, title, tags, account, name, u_id, user_link, img))
             con.commit()
-
-            print('success!')
-            time.sleep(0.5)
             con.close()
 
         elif i_tpye == '擦边球':
-            con = sqlite3.connect(path_setu + 'nearR18.db')
+            if os.path.exists('ATRI/sqlite/setu/nearR18.db'):
+                print('数据文件存在！')
+            else:
+                await session.send('数据库不存在，将在3秒后开始构建...')
+                time.sleep(3)
+                await session.send('开始构建数据库！')
+                con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'nearR18.db')
+                cur = con.cursor()
+                cur.execute('CREATE TABLE nearR18(pid PID, title TITLE, tags TAGS, account ACCOUNT, name NAME, u_id UID, user_link USERLINK, img IMG, UNIQUE(pid, title, tags, account, name, u_id, user_link, img))')
+                con.commit()
+                cur.close()
+                con.close()
+                await session.send('完成')
+ 
+            con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'nearR18.db')
             cur = con.cursor()
-            
             cur.execute('INSERT INTO nearR18 VALUES ("%s", "%s", "%s", "%s", "%s", "%s", "%s", "%s")'%(pid, title, tags, account, name, u_id, user_link, img))
             con.commit()
-
-            print('success!')
-            time.sleep(0.5)
             con.close()
 
         elif i_tpye == 'r18':
-            con = sqlite3.connect(path_setu + 'r18.db')
+            if os.path.exists('ATRI/sqlite/cloudmusic/cloudmusic.db'):
+                print('数据文件存在！')
+            else:
+                await session.send('数据库不存在，将在3秒后开始构建...')
+                time.sleep(3)
+                await session.send('开始构建数据库！')
+                con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'r18.db')
+                cur = con.cursor()
+                cur.execute('CREATE TABLE r18(pid PID, title TITLE, tags TAGS, account ACCOUNT, name NAME, u_id UID, user_link USERLINK, img IMG, UNIQUE(pid, title, tags, account, name, u_id, user_link, img))')
+                con.commit()
+                cur.close()
+                con.close()
+                await session.send('完成')
+
+            con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'r18.db')
             cur = con.cursor()
-            
             cur.execute('INSERT INTO r18 VALUES ("%s", "%s", "%s", "%s", "%s", "%s", "%s", "%s")'%(pid, title, tags, account, name, u_id, user_link, img))
             con.commit()
-
-            print('success!')
-            time.sleep(0.5)
             con.close()
         
-        await session.send('数据上传完成！')
+        end = time.perf_counter()
+        
+        await session.send(f'数据上传完成！\n耗时: {round(end - start, 3)}s')
 
 @on_command('upload_cloudmusic', aliases = ['上传网抑语', '网抑云', '网易云'], only_to_me = False)
 async def _(session: CommandSession):
     if session.event.user_id in master:
+        start = time.perf_counter()
         msg = session.event.raw_message.split(' ', 1)
         msg = msg[1]
 
-        con = sqlite3.connect(path_cloudmusic + 'cloudmusic.db')
+        if os.path.exists('ATRI/sqlite/cloudmusic/cloudmusic.db'):
+                print('数据文件存在！')
+        else:
+            await session.send('数据库不存在，将在3秒后开始构建...')
+            time.sleep(3)
+            await session.send('开始构建数据库！')
+            con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'cloudmusic' / 'cloudmusic.db')
+            cur = con.cursor()
+            cur.execute('CREATE TABLE cloudmusic(msg MSG, UNIQUE(msg))')
+            con.commit()
+            cur.close()
+            con.close()
+            await session.send('完成')
+ 
+        con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'cloudmusic' / 'cloudmusic.db')
         cur = con.cursor()
-
         cur.execute('INSERT INTO cloudmusic VALUES ("%s")'%(msg))
         con.commit()
-
-        print('success!')
-        time.sleep(0.5)
         con.close()
 
-        await session.send('数据上传完成！')
+        end = time.perf_counter()
+
+        await session.send(f'数据上传完成！\n耗时: {round(end - start, 3)}s')
 
 
 @on_command('data_list', aliases = ['查看上传数据', '数据总量'], only_to_me = False)
 async def _(session: CommandSession):
-    con = sqlite3.connect(path_setu + 'normal.db') # setu-normal
+    con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'normal.db') # setu-normal
     cur = con.cursor()
     cur.execute("select * from normal")
     data_normal = len(cur.fetchall())
     con.close()
 
-    con = sqlite3.connect(path_setu + 'nearR18.db') # setu-nearR18
+    con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'nearR18.db') # setu-nearR18
     cur = con.cursor()
     cur.execute("select * from nearR18")
     data_nearR18 = len(cur.fetchall())
     con.close()
 
-    con = sqlite3.connect(path_setu + 'r18.db') # setu-r18
+    con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'setu' / 'r18.db') # setu-r18
     cur = con.cursor()
     cur.execute("select * from r18")
     data_r18 = len(cur.fetchall())
     con.close()
 
-    con = sqlite3.connect(path_cloudmusic + 'cloudmusic.db') # cloudmusic
+    con = sqlite3.connect(Path('.') / 'ATRI' / 'sqlite' / 'cloudmusic' / 'cloudmusic.db') # cloudmusic
     cur = con.cursor()
     cur.execute("select * from cloudmusic")
     data_cloudmusic = len(cur.fetchall())