From 1c58b6e665b30c99846f5acd335195b56336f199 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B4=94=E4=B8=BA=E4=B9=8B?= <560397@gree.com.cn>
Date: Fri, 3 Nov 2023 17:59:07 +0800
Subject: [PATCH] Update project

---
 application/config/config.yaml            |  6 ++-
 application/extensions/init_sqlalchemy.py | 16 +++++-
 application/utils/__init__.py             |  2 +-
 application/utils/dsn/__init__.py         |  2 +-
 application/utils/dsn/dsn.py              | 61 +++++++++++++++++++----
 5 files changed, 70 insertions(+), 17 deletions(-)

diff --git a/application/config/config.yaml b/application/config/config.yaml
index 86fc09f..7f867b4 100644
--- a/application/config/config.yaml
+++ b/application/config/config.yaml
@@ -1,8 +1,10 @@
 System:
   Env: public
 
-MySQL:
-  User: root
+Database:
+  Type: mysql
+  Driver: pymysql
+  Username: root
   Password: localhost123
   Host: localhost
   Port: 3306
diff --git a/application/extensions/init_sqlalchemy.py b/application/extensions/init_sqlalchemy.py
index 1e47102..adc777e 100644
--- a/application/extensions/init_sqlalchemy.py
+++ b/application/extensions/init_sqlalchemy.py
@@ -13,7 +13,7 @@
 from flask import Flask
 from flask_sqlalchemy import SQLAlchemy
 
-from application.utils import dsn
+from application.utils import DatabaseURI
 
 db = SQLAlchemy()
 
@@ -25,5 +25,17 @@ def init_database(app: Flask) -> None:
     :param app: flask.Flask application instance
     :return: None
     """
-    app.config.setdefault('SQLALCHEMY_DATABASE_URI', dsn(app))
+    cfg = app.config.get('Database')
+    if cfg is None:
+        raise KeyError('Key Database error')
+    uri = DatabaseURI(
+        db_type=cfg.Type,
+        username=cfg.Username,
+        password=cfg.Password,
+        host=cfg.Host,
+        port=cfg.Port,
+        db=cfg.DB,
+        driver=cfg.Driver,
+    )
+    app.config.setdefault('SQLALCHEMY_DATABASE_URI', uri.create())
     db.init_app(app)
diff --git a/application/utils/__init__.py b/application/utils/__init__.py
index eb4d1d2..796c9b1 100644
--- a/application/utils/__init__.py
+++ b/application/utils/__init__.py
@@ -10,5 +10,5 @@
 # @Description :
 """
 
-from .dsn import dsn
+from .dsn import dsn, DatabaseURI
 from .elasticsearch import ElasticsearchUtils
diff --git a/application/utils/dsn/__init__.py b/application/utils/dsn/__init__.py
index 5711662..de7c835 100644
--- a/application/utils/dsn/__init__.py
+++ b/application/utils/dsn/__init__.py
@@ -10,4 +10,4 @@
 # @Description :
 """
 
-from .dsn import dsn
+from .dsn import DatabaseURI
diff --git a/application/utils/dsn/dsn.py b/application/utils/dsn/dsn.py
index 87937cf..13cc031 100644
--- a/application/utils/dsn/dsn.py
+++ b/application/utils/dsn/dsn.py
@@ -10,18 +10,57 @@
 # @Description :
 """
 
-from flask import Flask
 
-
-def dsn(app: Flask) -> str:
+class DatabaseURI:
     """
-    Initialize the MySQL dsn extension
+    The DatabaseURI class is a utility for generating SQLAlchemy database
+    connection strings for various types of databases.
 
-    :param app: flask.Flask application instance
-    :return: dsn
+    :param db_type: The type of the database. It could be "mysql", "postgresql", "oracle", "sqlite".
+    :param username: The username of the database. Not required for "sqlite".
+    :param password: The password of the database. Not required for "sqlite".
+    :param host: The host of the database. Not required for "sqlite".
+    :param port: The port of the database. Not required for "sqlite".
+    :param db: The database name. For "sqlite", it is the path to the database file.
+    :param driver: The python library to connect to the database. Not required but could be provided.
     """
-    mysql_cfg = app.config.get('MySQL')
-    if mysql_cfg is None:
-        raise KeyError('Key MySQL error')
-    return (f'mysql+pymysql://{app.config.MySQL.User}:'
-            f'{app.config.MySQL.Password}@{app.config.MySQL.Host}:{app.config.MySQL.Port}/{app.config.MySQL.DB}')
+
+    def __init__(self,
+                 db_type: str,
+                 username: str = None,
+                 password: str = None,
+                 host: str = None,
+                 port: str = None,
+                 db: str = None,
+                 driver: str = None
+                 ) -> None:
+        self.db_type = db_type
+        self.username = username
+        self.password = password
+        self.host = host
+        self.port = port
+        self.db = db
+        self.driver = driver
+
+    def create(self) -> str:
+        """
+        Creates a SQLAlchemy engine string for the specified database.
+
+        :return: A SQLAlchemy engine string.
+        """
+        # SQLite database only require a path to the database file
+        if self.db_type == 'sqlite':
+            return f'{self.db_type}:///{self.db}'
+        else:
+            # Other databases require a full connection string with optional driver
+            if self.driver:
+                return (f'{self.db_type}+{self.driver}://'
+                        f'{self.username}:{self.password}@{self.host}:{self.port}/{self.db}')
+            else:
+                return f'{self.db_type}://{self.username}:{self.password}@{self.host}:{self.port}/{self.db}'
+
+    def __repr__(self):
+        return f"<DatabaseURI({self.db_type})>"
+
+    def __str__(self):
+        return self.create()
-- 
GitLab