mirror of
https://github.com/rachelos/we-mp-rss.git
synced 2026-08-31 02:19:14 +08:00
t@123654
229 lines
11 KiB
Python
229 lines
11 KiB
Python
import os
|
|
import importlib
|
|
from typing import Dict, Type
|
|
from sqlalchemy import create_engine, MetaData, inspect
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
import logging
|
|
|
|
class DatabaseSynchronizer:
|
|
"""数据库模型同步器"""
|
|
|
|
def __init__(self, db_url: str, models_dir: str = "core/models"):
|
|
"""
|
|
初始化同步器
|
|
|
|
:param db_url: 数据库连接URL
|
|
:param models_dir: 模型目录路径
|
|
"""
|
|
self.db_url = db_url
|
|
self.models_dir = models_dir
|
|
self.engine = None
|
|
self.models = {}
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO)
|
|
self.logger = logging.getLogger("Sync")
|
|
|
|
def load_models(self) -> Dict[str, Type[declarative_base()]]:
|
|
"""动态加载所有模型类"""
|
|
self.models = {}
|
|
for filename in os.listdir(self.models_dir):
|
|
if filename.endswith(".py") and not filename.startswith("__"):
|
|
module_name = filename[:-3]
|
|
try:
|
|
module = importlib.import_module(f"core.models.{module_name}")
|
|
for name, obj in module.__dict__.items():
|
|
if isinstance(obj, type) and hasattr(obj, "__tablename__"):
|
|
self.models[obj.__tablename__] = obj
|
|
self.logger.info(f"成功加载模型模块: {module_name}")
|
|
except ImportError as e:
|
|
self.logger.warning(f"无法加载模型模块 {module_name}: {e}")
|
|
return self.models
|
|
|
|
def _map_types_for_database(self, model):
|
|
"""为不同数据库处理特殊类型映射"""
|
|
for column in model.__table__.columns:
|
|
type_str = str(column.type).upper()
|
|
|
|
# SQLite类型映射
|
|
if "sqlite" in self.db_url:
|
|
# 检查多种可能的MEDIUMTEXT表示形式
|
|
if (hasattr(column.type, "__visit_name__") and column.type.__visit_name__ == "MEDIUMTEXT") or \
|
|
"MEDIUMTEXT" in type_str or \
|
|
getattr(column.type, "__class__", None).__name__ == "MEDIUMTEXT":
|
|
from sqlalchemy import Text
|
|
column.type = Text()
|
|
self.logger.debug(f"已将列 {column.name} 的类型从 MEDIUMTEXT 映射为 Text")
|
|
|
|
# PostgreSQL类型映射
|
|
elif "postgresql" in self.db_url or "postgres" in self.db_url:
|
|
# MEDIUMTEXT映射为TEXT
|
|
if (hasattr(column.type, "__visit_name__") and column.type.__visit_name__ == "MEDIUMTEXT") or \
|
|
"MEDIUMTEXT" in type_str or \
|
|
getattr(column.type, "__class__", None).__name__ == "MEDIUMTEXT":
|
|
from sqlalchemy import Text
|
|
column.type = Text()
|
|
self.logger.debug(f"已将列 {column.name} 的类型从 MEDIUMTEXT 映射为 Text")
|
|
|
|
# LONGTEXT映射为TEXT
|
|
if "LONGTEXT" in type_str or \
|
|
getattr(column.type, "__class__", None).__name__ == "LONGTEXT":
|
|
from sqlalchemy import Text
|
|
column.type = Text()
|
|
self.logger.debug(f"已将列 {column.name} 的类型从 LONGTEXT 映射为 Text")
|
|
|
|
# TINYINT映射为SMALLINT
|
|
if "TINYINT" in type_str or \
|
|
getattr(column.type, "__class__", None).__name__ == "TINYINT":
|
|
from sqlalchemy import SmallInteger
|
|
column.type = SmallInteger()
|
|
self.logger.debug(f"已将列 {column.name} 的类型从 TINYINT 映射为 SmallInteger")
|
|
|
|
def _check_database_permissions(self):
|
|
"""检查数据库权限"""
|
|
try:
|
|
with self.engine.begin() as conn:
|
|
# 检查是否可以创建表
|
|
if "postgresql" in self.db_url or "postgres" in self.db_url:
|
|
# 检查当前用户权限
|
|
result = conn.execute("SELECT current_user, current_database(), current_schema()")
|
|
user_info = result.fetchone()
|
|
self.logger.info(f"当前用户: {user_info[0]}, 数据库: {user_info[1]}, Schema: {user_info[2]}")
|
|
|
|
# 检查schema权限
|
|
result = conn.execute("""
|
|
SELECT has_schema_privilege(current_user, 'public', 'CREATE') as can_create,
|
|
has_schema_privilege(current_user, 'public', 'USAGE') as can_use
|
|
""")
|
|
perms = result.fetchone()
|
|
|
|
if not perms[0]: # 没有CREATE权限
|
|
self.logger.error("当前用户没有在public schema中创建表的权限")
|
|
self.logger.info("请联系数据库管理员执行以下命令:")
|
|
self.logger.info(f"GRANT CREATE ON SCHEMA public TO {user_info[0]};")
|
|
return False
|
|
|
|
if not perms[1]: # 没有USAGE权限
|
|
self.logger.error("当前用户没有使用public schema的权限")
|
|
self.logger.info("请联系数据库管理员执行以下命令:")
|
|
self.logger.info(f"GRANT USAGE ON SCHEMA public TO {user_info[0]};")
|
|
return False
|
|
|
|
return True
|
|
except Exception as e:
|
|
self.logger.warning(f"权限检查失败: {e}")
|
|
return True # 如果检查失败，继续尝试
|
|
|
|
def sync(self):
|
|
"""同步模型到数据库"""
|
|
try:
|
|
self.engine = create_engine(self.db_url)
|
|
|
|
# 检查数据库权限
|
|
if not self._check_database_permissions():
|
|
return False
|
|
|
|
metadata = MetaData()
|
|
|
|
# 反射现有数据库结构
|
|
metadata.reflect(bind=self.engine)
|
|
|
|
# 处理不同数据库的特殊类型映射
|
|
for model in self.models.values():
|
|
self._map_types_for_database(model)
|
|
|
|
# 加载模型
|
|
if not self.models:
|
|
self.load_models()
|
|
if not self.models:
|
|
self.logger.error("没有找到任何模型类")
|
|
return False
|
|
|
|
# 为不同数据库类型处理自增主键
|
|
if "sqlite" in self.db_url:
|
|
# SQLite使用AUTOINCREMENT
|
|
pass # SQLAlchemy默认处理
|
|
elif "mysql" in self.db_url:
|
|
# MySQL使用AUTO_INCREMENT
|
|
pass # SQLAlchemy默认处理
|
|
elif "postgresql" in self.db_url or "postgres" in self.db_url:
|
|
# PostgreSQL使用SERIAL或IDENTITY
|
|
pass # SQLAlchemy默认处理
|
|
|
|
# 创建或更新表结构
|
|
for model in self.models.values():
|
|
table_name = model.__tablename__
|
|
inspector = inspect(self.engine)
|
|
|
|
try:
|
|
if not inspector.has_table(table_name):
|
|
# 尝试创建表
|
|
model.metadata.create_all(self.engine)
|
|
self.logger.info(f"创建表: {table_name}")
|
|
else:
|
|
# 检查字段差异并更新表
|
|
existing_columns = {c["name"]: c for c in inspector.get_columns(table_name)}
|
|
model_columns = {c.name: c for c in model.__table__.columns}
|
|
|
|
# 检查新增或修改的字段
|
|
for col_name, model_col in model_columns.items():
|
|
if col_name not in existing_columns:
|
|
# 新增字段 - 根据数据库类型调整语法
|
|
from sqlalchemy import text
|
|
try:
|
|
with self.engine.begin() as conn:
|
|
if "postgresql" in self.db_url or "postgres" in self.db_url:
|
|
# PostgreSQL语法
|
|
conn.execute(text(f'ALTER TABLE "{table_name}" ADD COLUMN "{col_name}" {model_col.type}'))
|
|
else:
|
|
# SQLite和MySQL语法
|
|
conn.execute(text(f"ALTER TABLE {table_name} ADD COLUMN {col_name} {model_col.type}"))
|
|
self.logger.info(f"新增字段: {table_name}.{col_name}")
|
|
except SQLAlchemyError as e:
|
|
self.logger.error(f"添加字段 {table_name}.{col_name} 失败: {e}")
|
|
|
|
self.logger.info(f"表已同步: {table_name}")
|
|
|
|
except SQLAlchemyError as e:
|
|
self.logger.error(f"处理表 {table_name} 时出错: {e}")
|
|
if "permission denied" in str(e).lower():
|
|
self.logger.error("权限不足，请检查数据库用户权限")
|
|
return False
|
|
continue
|
|
|
|
self.logger.info("模型同步完成")
|
|
return True
|
|
except SQLAlchemyError as e:
|
|
self.logger.error(f"数据库同步失败: {e}")
|
|
if "permission denied" in str(e).lower():
|
|
self.logger.error("数据库权限不足，请检查以下几点:")
|
|
self.logger.error("1. 确保数据库用户有CREATE权限")
|
|
self.logger.error("2. 确保数据库用户有USAGE权限")
|
|
self.logger.error("3. 如果是PostgreSQL，请联系管理员执行权限授予命令")
|
|
return False
|
|
except Exception as e:
|
|
self.logger.error(f"同步过程中发生未知错误: {e}")
|
|
return False
|
|
finally:
|
|
if self.engine:
|
|
self.engine.dispose()
|
|
|
|
def main():
|
|
# 示例使用 - 支持多种数据库
|
|
# SQLite
|
|
# synchronizer = DatabaseSynchronizer(db_url="sqlite:///data/db.db")
|
|
|
|
# PostgreSQL
|
|
# synchronizer = DatabaseSynchronizer(db_url="postgresql://username:password@localhost:5432/dbname")
|
|
|
|
# MySQL
|
|
# synchronizer = DatabaseSynchronizer(db_url="mysql+pymysql://username:password@localhost:3306/dbname")
|
|
from core.config import cfg
|
|
db_url=cfg.get("db","sqlite:///data/db.db")
|
|
synchronizer = DatabaseSynchronizer(db_url=db_url)
|
|
synchronizer.sync()
|
|
|
|
if __name__ == "__main__":
|
|
main() |