diff --git a/src/db/__init__.py b/src/db/__init__.py index 74c7bcf..3e39539 100644 --- a/src/db/__init__.py +++ b/src/db/__init__.py @@ -6,12 +6,47 @@ # @File Name: __init__.py.py +import os from src.config import Config -if Config.database == 'sqlite': +database = Config.database + + +def split_(origin: str, type_: int = 0) -> dict: + """ + + :param origin: + :param type_: 0 -> Mysql + :return: + """ + origin = origin.split('|')[-1] + host = origin.split('@')[-1].split('/')[0].split(':')[0] + port = origin.split('@')[-1].split('/')[0].split(':')[1] + if type_ == 0: # Mysql + database_ = origin.split('@')[-1].split('/')[-1] + user = origin.split('@')[0].split(':')[0] + password = origin.split('@')[0].split(':')[1] + return { + 'host': host, + 'port': port, + 'database': database_, + 'user': user, + 'password': password + } + + +if database == 'sqlite': + import sqlite3 as operator + from src.utils.t_download import download + + if not os.path.exists('./src/db/data.sqlite'): + download('https://markusjoe.github.io/static_file_hosting/static/counter/data.sqlite') from src.db.db import SQLite as Database else: + import pymysql as operator + + for k, v in zip(split_(origin=database, type_=0)): + os.environ[k] = v from src.db.db import MySQL as Database - -__all__ = [Database] +__all__ = [Database, operator] diff --git a/src/db/db.py b/src/db/db.py index 51ca29c..e70e690 100644 --- a/src/db/db.py +++ b/src/db/db.py @@ -7,12 +7,10 @@ import os -import sqlite3 -from src.config import Config -from src.utils.t_download import download +from src.db import operator -class Database: +class DatabaseSQL: def __init__(self): self.conn = None self.cursor = None @@ -51,14 +49,14 @@ class Database: return result -class SQLite(Database): +class SQLite(DatabaseSQL): def __init__(self): super().__init__() - self.conn = sqlite3.connect('./src/db/data.sqlite') + self.conn = operator.connect('./src/db/data.sqlite') self.cursor = self.conn.cursor() -class MySQL(Database): +class MySQL(DatabaseSQL): def __init__(self): super().__init__() user = os.getenv('m_user') @@ -66,23 +64,10 @@ class MySQL(Database): host = os.getenv('m_host') port = int(os.getenv('m_port')) db = os.getenv('m_database') - self.conn = pymysql.connect(host=host, - user=user, - password=pwd, - port=port, - database=db - ) + self.conn = operator.connect(host=host, + user=user, + password=pwd, + port=port, + database=db + ) self.cursor = self.conn.cursor() - - -if __name__ != '__main__': - if Config.database == 'sqlite': - print('Using SQLite Database') - if not os.path.exists('./src/db/data.sqlite'): - print('database file not exists, start downloading...') - download('https://pac.rtst.tech/static_file_hosting/static/counter/data.sqlite') - print('Done.') - else: - print('Using MySQL Database') - if not os.getenv('m_user'): - import pymysql