From 2822b61d8e1e620f79cb972d6cc5fc0405885380 Mon Sep 17 00:00:00 2001
From: MarkusJoe <markushammered@gmail.com>
Date: Wed, 14 Sep 2022 12:46:06 +0800
Subject: [PATCH] =?UTF-8?q?update:=20=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?=
 =?UTF-8?q?=E4=BB=A5=E4=BE=BF=E6=9B=B4=E6=96=B9=E4=BE=BF=E7=9A=84=E5=BC=80?=
 =?UTF-8?q?=E5=8F=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/db/__init__.py | 41 ++++++++++++++++++++++++++++++++++++++---
 src/db/db.py       | 37 +++++++++++--------------------------
 2 files changed, 49 insertions(+), 29 deletions(-)

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