diff --git a/src/db/__init__.py b/src/db/__init__.py index 676a754..60ea2d8 100644 --- a/src/db/__init__.py +++ b/src/db/__init__.py @@ -7,17 +7,21 @@ import os -import sqlite3 as operator from urllib.request import urlretrieve from src.config import Config -from src.db.db import SQLite as Database -database = Config.database +database = Config.database # operator -if not os.path.exists('./src/db/data.sqlite'): +if not os.path.exists('./src/db/data.sqlite') and \ + Config.database == 'sqlite3': print('Downloading database file. Please wait...') file_url = 'https://static.rtast.cn/data.sqlite' - urlretrieve(file_url, './src/db/data.sqlite') + urlretrieve(file_url, './src/db/data.sqlite') # standard lib for downloading file print('Download database file successfully.') -__all__ = [Database, operator] +if Config.database == 'sqlite3': + from src.db.db import SQLite as Database + +else: + from src.db.db import MySQL as Database +__all__ = [Database] diff --git a/src/db/db.py b/src/db/db.py index 75ed72a..3ff90c6 100644 --- a/src/db/db.py +++ b/src/db/db.py @@ -6,7 +6,12 @@ # @File Name: db.py -from src.db import operator +from src.config import Config + +if Config.database == 'sqlite3': + import sqlite3 as operator +else: + import pymysql as operator class BaseSQL: @@ -29,7 +34,7 @@ class BaseSQL: return result def insert(self, _id: str) -> bool: - self.cursor.execute('insert into data (id, times) values (%s, 1);' % _id) + self.cursor.execute('insert into data (id, times) values ("%s", 1);' % _id) return True def update(self, _id: str, times: int) -> bool: @@ -53,3 +58,21 @@ class SQLite(BaseSQL): super().__init__() self.conn = operator.connect('./src/db/data.sqlite') self.cursor = self.conn.cursor() + + +class MySQL(BaseSQL): + def __init__(self): + super().__init__() + _CONFIG = Config.database.split('@') + user = _CONFIG[0].split(':')[0] + pwd = _CONFIG[0].split(':')[1] + host = _CONFIG[1].split(':')[0] + port = int(_CONFIG[1].split(':')[1].split('/')[0]) + db = _CONFIG[1].split('/')[1] + + self.conn = operator.connect(user=user, + passwd=pwd, + host=host, + port=port, + database=db) + self.cursor = self.conn.cursor()