185 lines
7.4 KiB
Python
185 lines
7.4 KiB
Python
"""
|
|
数据库连接与初始化
|
|
"""
|
|
import os
|
|
import pymysql
|
|
from werkzeug.security import generate_password_hash
|
|
|
|
try:
|
|
from config import mysql_host as config_mysql_host
|
|
from config import mysql_user as config_mysql_user
|
|
from config import mysql_password as config_mysql_password
|
|
from config import mysql_database as config_mysql_database
|
|
from config import mysql_host_source as config_mysql_host_source
|
|
from config import mysql_user_source as config_mysql_user_source
|
|
from config import mysql_database_source as config_mysql_database_source
|
|
except ImportError:
|
|
config_mysql_host = 'localhost'
|
|
config_mysql_user = 'root'
|
|
config_mysql_password = ''
|
|
config_mysql_database = 'maixiang_ai'
|
|
config_mysql_host_source = 'fallback.default'
|
|
config_mysql_user_source = 'fallback.default'
|
|
config_mysql_database_source = 'fallback.default'
|
|
|
|
mysql_host = config_mysql_host
|
|
mysql_user = config_mysql_user
|
|
mysql_password = config_mysql_password
|
|
mysql_database = config_mysql_database
|
|
mysql_host_source = config_mysql_host_source
|
|
mysql_user_source = config_mysql_user_source
|
|
mysql_database_source = config_mysql_database_source
|
|
|
|
|
|
def describe_db_target():
|
|
return (
|
|
f"{mysql_user}@{mysql_host}/{mysql_database} "
|
|
f"(host={mysql_host_source}, user={mysql_user_source}, database={mysql_database_source})"
|
|
)
|
|
|
|
|
|
def get_db():
|
|
return pymysql.connect(
|
|
host=mysql_host,
|
|
user=mysql_user,
|
|
password=mysql_password,
|
|
database=mysql_database,
|
|
charset='utf8mb4',
|
|
cursorclass=pymysql.cursors.DictCursor
|
|
)
|
|
|
|
|
|
def _create_initial_admin():
|
|
"""若没有任何管理员，则创建默认超级管理员（首次启动时，仅一个）"""
|
|
try:
|
|
conn = get_db()
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT id FROM users WHERE role = 'super_admin' LIMIT 1")
|
|
if cur.fetchone():
|
|
conn.close()
|
|
return
|
|
admin_user = os.environ.get('ADMIN_USER', 'admin')
|
|
admin_pwd = os.environ.get('ADMIN_PASSWORD', 'admin123')
|
|
pwd_hash = generate_password_hash(admin_pwd, method='pbkdf2:sha256')
|
|
cur.execute(
|
|
"INSERT INTO users (username, password_hash, is_admin, role) VALUES (%s, %s, 1, 'super_admin')",
|
|
(admin_user, pwd_hash)
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def init_db():
|
|
"""初始化数据库表，若不存在则创建"""
|
|
conn = pymysql.connect(
|
|
host=mysql_host,
|
|
user=mysql_user,
|
|
password=mysql_password,
|
|
charset='utf8mb4'
|
|
)
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute(f"CREATE DATABASE IF NOT EXISTS `{mysql_database}` DEFAULT CHARSET utf8mb4")
|
|
cur.execute(f"USE `{mysql_database}`")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
username VARCHAR(64) NOT NULL UNIQUE,
|
|
password_hash VARCHAR(256) NOT NULL,
|
|
is_admin TINYINT(1) DEFAULT 0,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS image_history (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
user_id INT NOT NULL,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
panel_type VARCHAR(64) DEFAULT '',
|
|
original_urls JSON,
|
|
params JSON,
|
|
result_urls JSON,
|
|
long_image_url VARCHAR(1024) DEFAULT NULL,
|
|
INDEX idx_user_created (user_id, created_at DESC)
|
|
)
|
|
""")
|
|
try:
|
|
cur.execute("ALTER TABLE image_history ADD COLUMN long_image_url VARCHAR(1024) DEFAULT NULL")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("ALTER TABLE users ADD COLUMN machine VARCHAR(64) DEFAULT NULL")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("ALTER TABLE users ADD COLUMN role VARCHAR(20) DEFAULT 'normal'")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("ALTER TABLE users ADD COLUMN created_by_id INT NULL")
|
|
except Exception:
|
|
pass
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS web_config (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
version VARCHAR(64) NOT NULL,
|
|
file_url VARCHAR(1024) NOT NULL,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS columns (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
name VARCHAR(128) NOT NULL COMMENT '栏目名',
|
|
column_key VARCHAR(64) NOT NULL COMMENT '栏目标识',
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE KEY uk_column_key (column_key)
|
|
)
|
|
""")
|
|
try:
|
|
cur.execute("ALTER TABLE columns ADD COLUMN menu_type VARCHAR(20) NOT NULL DEFAULT 'app' COMMENT '菜单类型: app/admin' AFTER column_key")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("ALTER TABLE columns ADD COLUMN route_path VARCHAR(255) NOT NULL DEFAULT '' COMMENT '菜单路由或页面标识' AFTER menu_type")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("ALTER TABLE columns ADD COLUMN sort_order INT NOT NULL DEFAULT 0 COMMENT '菜单排序' AFTER route_path")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("UPDATE columns SET menu_type = 'app' WHERE menu_type IS NULL OR menu_type = ''")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
cur.execute("UPDATE columns SET sort_order = id WHERE sort_order IS NULL OR sort_order = 0")
|
|
except Exception:
|
|
pass
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS user_column_permission (
|
|
user_id INT NOT NULL,
|
|
column_id INT NOT NULL,
|
|
PRIMARY KEY (user_id, column_id),
|
|
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
|
|
FOREIGN KEY (column_id) REFERENCES columns(id) ON DELETE CASCADE
|
|
)
|
|
""")
|
|
try:
|
|
cur.execute("UPDATE users SET role = 'normal' WHERE (role IS NULL OR role = '') AND (is_admin = 0 OR is_admin IS NULL)")
|
|
cur.execute("SELECT MIN(id) AS mid FROM users WHERE is_admin = 1")
|
|
row = cur.fetchone()
|
|
if row and row.get('mid'):
|
|
mid = row['mid']
|
|
cur.execute("UPDATE users SET role = 'super_admin' WHERE id = %s", (mid,))
|
|
cur.execute("UPDATE users SET role = 'admin' WHERE is_admin = 1 AND id != %s", (mid,))
|
|
cur.execute("UPDATE users SET created_by_id = %s WHERE role = 'admin' AND (created_by_id IS NULL)", (mid,))
|
|
except Exception:
|
|
pass
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
_create_initial_admin()
|