269 lines
9.1 KiB
Python
269 lines
9.1 KiB
Python
|
"""数据库初始化：Doris / MySQL / SQLite 建表与迁移"""
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
|
|||
|
|
from sqlalchemy import event, inspect as sa_inspect
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
|
|||
|
|
from models import db
|
|||
|
|
|
|||
|
|
_doris_id_listener_registered = False
|
|||
|
|
|
|||
|
|
# 建表顺序（逻辑依赖顺序，Doris 不强制外键）
|
|||
|
|
TABLE_ORDER = [
|
|||
|
|
'roles', 'users', 'teachers', 'students', 'courses', 'classes',
|
|||
|
|
'class_students', 'student_accounts', 'recharge_activities',
|
|||
|
|
'recharge_records', 'consumption_records', 'refund_records',
|
|||
|
|
'transfer_records', 'stop_records', 'attendance_records',
|
|||
|
|
'operation_logs', 'schedules', 'monthly_snapshots', 'class_records',
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
DORIS_BUCKETS = int(os.environ.get('DORIS_BUCKETS', '1'))
|
|||
|
|
DORIS_REPLICATION_NUM = os.environ.get('DORIS_REPLICATION_NUM', '1')
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_uri():
|
|||
|
|
from flask import current_app
|
|||
|
|
return current_app.config.get('SQLALCHEMY_DATABASE_URI', '')
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_doris():
|
|||
|
|
"""是否连接 Apache Doris（MySQL 协议，FE 默认 9030）"""
|
|||
|
|
if os.environ.get('DB_TYPE', '').lower() == 'doris':
|
|||
|
|
return True
|
|||
|
|
uri = _get_uri()
|
|||
|
|
if 'doris' in uri.lower():
|
|||
|
|
return True
|
|||
|
|
m = re.search(r':(\d+)/', uri)
|
|||
|
|
return bool(m and m.group(1) == '9030')
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_mysql():
|
|||
|
|
"""是否 MySQL 协议（含 Doris）"""
|
|||
|
|
return _get_uri().startswith('mysql')
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_sqlite():
|
|||
|
|
return _get_uri().startswith('sqlite')
|
|||
|
|
|
|||
|
|
|
|||
|
|
def quote_ident(name):
|
|||
|
|
return f'`{name}`'
|
|||
|
|
|
|||
|
|
|
|||
|
|
def doris_next_id(session, table_name, pk_name='id'):
|
|||
|
|
"""Doris 无 LAST_INSERT_ID，用 MAX(id)+1 分配主键"""
|
|||
|
|
sql = f'SELECT COALESCE(MAX({quote_ident(pk_name)}), 0) + 1 FROM {quote_ident(table_name)}'
|
|||
|
|
return int(session.execute(db.text(sql)).scalar() or 1)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def register_doris_id_listener():
|
|||
|
|
"""插入前为无主键记录分配 id（仅 Doris）"""
|
|||
|
|
global _doris_id_listener_registered
|
|||
|
|
if _doris_id_listener_registered:
|
|||
|
|
return
|
|||
|
|
_doris_id_listener_registered = True
|
|||
|
|
|
|||
|
|
@event.listens_for(Session, 'before_flush')
|
|||
|
|
def _assign_doris_ids(session, flush_context, instances):
|
|||
|
|
if not is_doris():
|
|||
|
|
return
|
|||
|
|
for obj in session.new:
|
|||
|
|
mapper = sa_inspect(obj.__class__)
|
|||
|
|
if not mapper.tables:
|
|||
|
|
continue
|
|||
|
|
table = mapper.tables[0]
|
|||
|
|
for col in mapper.primary_key:
|
|||
|
|
if not getattr(col, 'autoincrement', False):
|
|||
|
|
continue
|
|||
|
|
if getattr(obj, col.key) is not None:
|
|||
|
|
continue
|
|||
|
|
new_id = doris_next_id(session, table.name, col.key)
|
|||
|
|
setattr(obj, col.key, new_id)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _col_type_doris(col):
|
|||
|
|
"""SQLAlchemy 列类型 -> Doris DDL 类型"""
|
|||
|
|
if col.primary_key and getattr(col, 'autoincrement', False):
|
|||
|
|
# 主键由应用分配（Doris 不支持可靠的 AUTO_INCREMENT / LAST_INSERT_ID）
|
|||
|
|
return 'BIGINT NOT NULL'
|
|||
|
|
|
|||
|
|
type_str = str(col.type).upper()
|
|||
|
|
if 'BIGINT' in type_str:
|
|||
|
|
return 'BIGINT'
|
|||
|
|
if 'SMALLINT' in type_str or 'TINYINT' in type_str:
|
|||
|
|
return 'SMALLINT'
|
|||
|
|
if 'INTEGER' in type_str or type_str == 'INT':
|
|||
|
|
return 'INT'
|
|||
|
|
if 'VARCHAR' in type_str:
|
|||
|
|
return type_str.replace('NVARCHAR', 'VARCHAR')
|
|||
|
|
if 'TEXT' in type_str:
|
|||
|
|
return 'STRING'
|
|||
|
|
if 'DATETIME' in type_str:
|
|||
|
|
return 'DATETIME'
|
|||
|
|
if type_str == 'DATE' or type_str.startswith('DATE('):
|
|||
|
|
return 'DATE'
|
|||
|
|
if 'NUMERIC' in type_str or 'DECIMAL' in type_str:
|
|||
|
|
return type_str.replace('NUMERIC', 'DECIMAL')
|
|||
|
|
if 'FLOAT' in type_str or 'DOUBLE' in type_str:
|
|||
|
|
return type_str
|
|||
|
|
if 'BOOLEAN' in type_str:
|
|||
|
|
return 'TINYINT'
|
|||
|
|
return type_str
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _default_clause_doris(col):
|
|||
|
|
"""生成 Doris 列 DEFAULT 子句"""
|
|||
|
|
if col.primary_key and getattr(col, 'autoincrement', False):
|
|||
|
|
return ''
|
|||
|
|
if col.server_default is not None:
|
|||
|
|
return ''
|
|||
|
|
if col.default is None:
|
|||
|
|
return ''
|
|||
|
|
if not getattr(col.default, 'is_scalar', True):
|
|||
|
|
return ''
|
|||
|
|
val = col.default.arg
|
|||
|
|
if val is None:
|
|||
|
|
return ''
|
|||
|
|
if isinstance(val, bool):
|
|||
|
|
return f' DEFAULT "{int(val)}"'
|
|||
|
|
if isinstance(val, (int, float)):
|
|||
|
|
return f' DEFAULT "{val}"'
|
|||
|
|
if isinstance(val, str):
|
|||
|
|
return f" DEFAULT '{val.replace(chr(39), chr(39) + chr(39))}'"
|
|||
|
|
return ''
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _build_doris_column(col):
|
|||
|
|
parts = [quote_ident(col.name), _col_type_doris(col)]
|
|||
|
|
if not col.nullable and not (col.primary_key and getattr(col, 'autoincrement', False)):
|
|||
|
|
parts.append('NOT NULL')
|
|||
|
|
parts.append(_default_clause_doris(col))
|
|||
|
|
return ' '.join(p for p in parts if p)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_doris_create_sql(table):
|
|||
|
|
"""根据 SQLAlchemy Table 生成 Doris CREATE TABLE 语句"""
|
|||
|
|
pk_cols = [c for c in table.columns if c.primary_key]
|
|||
|
|
pk_col = pk_cols[0] if pk_cols else table.columns[0]
|
|||
|
|
pk_name = pk_col.name
|
|||
|
|
|
|||
|
|
col_defs = [_build_doris_column(c) for c in table.columns]
|
|||
|
|
tname = quote_ident(table.name)
|
|||
|
|
|
|||
|
|
return (
|
|||
|
|
f'CREATE TABLE IF NOT EXISTS {tname} (\n'
|
|||
|
|
f' {",\n ".join(col_defs)}\n'
|
|||
|
|
f') UNIQUE KEY({quote_ident(pk_name)})\n'
|
|||
|
|
f'DISTRIBUTED BY HASH({quote_ident(pk_name)}) BUCKETS {DORIS_BUCKETS}\n'
|
|||
|
|
f'PROPERTIES ("replication_num" = "{DORIS_REPLICATION_NUM}")'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _table_by_name(name):
|
|||
|
|
for mapper in db.Model.registry.mappers:
|
|||
|
|
if mapper.class_.__tablename__ == name:
|
|||
|
|
return mapper.class_.__table__
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def doris_create_all():
|
|||
|
|
"""Doris：按模型创建缺失的表（首次启动 / 新表）"""
|
|||
|
|
inspector = sa_inspect(db.engine)
|
|||
|
|
existing = set(inspector.get_table_names())
|
|||
|
|
|
|||
|
|
created = []
|
|||
|
|
for table_name in TABLE_ORDER:
|
|||
|
|
if table_name in existing:
|
|||
|
|
continue
|
|||
|
|
table = _table_by_name(table_name)
|
|||
|
|
if table is None:
|
|||
|
|
print(f' [doris_create] Skip unknown table: {table_name}')
|
|||
|
|
continue
|
|||
|
|
sql = build_doris_create_sql(table)
|
|||
|
|
try:
|
|||
|
|
db.session.execute(db.text(sql))
|
|||
|
|
db.session.commit()
|
|||
|
|
created.append(table_name)
|
|||
|
|
print(f' [doris_create] Created table: {table_name}')
|
|||
|
|
except Exception as e:
|
|||
|
|
db.session.rollback()
|
|||
|
|
print(f' [doris_create] Failed {table_name}: {e}')
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
# 处理 TABLE_ORDER 未列出的模型表
|
|||
|
|
for mapper in db.Model.registry.mappers:
|
|||
|
|
table = mapper.class_.__table__
|
|||
|
|
if table.name in existing or table.name in created:
|
|||
|
|
continue
|
|||
|
|
if table.name in TABLE_ORDER:
|
|||
|
|
continue
|
|||
|
|
sql = build_doris_create_sql(table)
|
|||
|
|
try:
|
|||
|
|
db.session.execute(db.text(sql))
|
|||
|
|
db.session.commit()
|
|||
|
|
print(f' [doris_create] Created table: {table.name}')
|
|||
|
|
except Exception as e:
|
|||
|
|
db.session.rollback()
|
|||
|
|
print(f' [doris_create] Failed {table.name}: {e}')
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
return created
|
|||
|
|
|
|||
|
|
|
|||
|
|
KESHI_LEDGER_COLUMNS = [
|
|||
|
|
('students', 'display_name', 'VARCHAR(50)'),
|
|||
|
|
('students', 'nickname', 'VARCHAR(50)'),
|
|||
|
|
('courses', 'course_code', 'VARCHAR(20)'),
|
|||
|
|
('student_accounts', 'unit_price', 'DECIMAL(10,4) DEFAULT "0"'),
|
|||
|
|
('student_accounts', 'original_price_per_lesson', 'DECIMAL(10,4)'),
|
|||
|
|
('student_accounts', 'account_status', 'VARCHAR(20) DEFAULT "active"'),
|
|||
|
|
('recharge_records', 'gift_clawed_back', 'SMALLINT DEFAULT "0"'),
|
|||
|
|
('refund_records', 'deduct_gifted_amount', 'DECIMAL(12,2) DEFAULT "0"'),
|
|||
|
|
('refund_records', 'promo_clawback_hours', 'DECIMAL(10,2) DEFAULT "0"'),
|
|||
|
|
('consumption_records', 'unit_price_at_consume', 'DECIMAL(10,4)'),
|
|||
|
|
('student_accounts', 'version', 'INT DEFAULT "0"'),
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def migrate_keshi_ledger_columns():
|
|||
|
|
"""课时核对表相关列：Doris 下确保 ALTER 执行（补齐 auto_migrate 漏迁场景）"""
|
|||
|
|
if not is_mysql():
|
|||
|
|
return
|
|||
|
|
for table_name, col_name, col_type in KESHI_LEDGER_COLUMNS:
|
|||
|
|
try:
|
|||
|
|
rows = db.session.execute(
|
|||
|
|
db.text(f'SHOW COLUMNS FROM {quote_ident(table_name)}')
|
|||
|
|
).fetchall()
|
|||
|
|
existing = {r[0] for r in rows}
|
|||
|
|
except Exception:
|
|||
|
|
continue
|
|||
|
|
if col_name in existing:
|
|||
|
|
continue
|
|||
|
|
sql = (
|
|||
|
|
f'ALTER TABLE {quote_ident(table_name)} '
|
|||
|
|
f'ADD COLUMN {quote_ident(col_name)} {col_type}'
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
db.session.execute(db.text(sql))
|
|||
|
|
db.session.commit()
|
|||
|
|
print(f' [migrate_keshi] Added {table_name}.{col_name}')
|
|||
|
|
except Exception as e:
|
|||
|
|
db.session.rollback()
|
|||
|
|
print(f' [migrate_keshi] Failed {table_name}.{col_name}: {e}')
|
|||
|
|
|
|||
|
|
|
|||
|
|
def bootstrap_database():
|
|||
|
|
"""统一入口：建表 + 自动迁移"""
|
|||
|
|
if is_doris():
|
|||
|
|
print('[bootstrap] Database: Doris')
|
|||
|
|
register_doris_id_listener()
|
|||
|
|
doris_create_all()
|
|||
|
|
else:
|
|||
|
|
print(f'[bootstrap] Database: {"MySQL" if is_mysql() else "SQLite"}')
|
|||
|
|
db.create_all()
|
|||
|
|
|
|||
|
|
from utils import auto_migrate
|
|||
|
|
auto_migrate()
|
|||
|
|
migrate_keshi_ledger_columns()
|