Files
geMoldInsight/src/database/init_db.py
T

257 lines
12 KiB
Python
Raw Normal View History

import asyncio
import sys
from pathlib import Path
from sqlalchemy import text
project_root = Path(__file__).parent.parent.parent
src_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(src_root))
from sqlalchemy import select
from database.database import db_manager
from models.database import User, Role, Permission, UserRole, RolePermission
from services.auth_service import get_password_hash
from config.settings import settings
from utils.logger import get_logger
logger = get_logger(__name__)
DEFAULT_PERMISSIONS = [
{"code": "view_dashboard", "name": "查看仪表盘", "module": "dashboard"},
{"code": "view_moldinsight", "name": "使用模具分析", "module": "moldinsight"},
{"code": "upload_file", "name": "上传文件", "module": "moldinsight"},
{"code": "view_history", "name": "查看历史记录", "module": "moldinsight"},
{"code": "view_inventory", "name": "查看库存", "module": "inventory"},
{"code": "manage_inventory", "name": "管理库存", "module": "inventory"},
{"code": "view_products", "name": "查看产品", "module": "inventory"},
{"code": "manage_products", "name": "管理产品", "module": "inventory"},
{"code": "view_suppliers", "name": "查看供应商", "module": "inventory"},
{"code": "manage_suppliers", "name": "管理供应商", "module": "inventory"},
{"code": "view_customers", "name": "查看客户", "module": "inventory"},
{"code": "manage_customers", "name": "管理客户", "module": "inventory"},
{"code": "view_finance", "name": "查看财务", "module": "finance"},
{"code": "manage_receipts", "name": "管理收款", "module": "finance"},
{"code": "manage_payments", "name": "管理付款", "module": "finance"},
{"code": "void_finance_transaction", "name": "作废财务单据", "module": "finance"},
{"code": "view_users", "name": "查看用户", "module": "admin"},
{"code": "manage_users", "name": "管理用户", "module": "admin"},
{"code": "manage_roles", "name": "管理角色", "module": "admin"},
]
DEFAULT_ROLES = [
{"code": "admin", "name": "管理员", "description": "系统管理员,拥有所有权限", "is_system": True, "permissions": ["view_dashboard", "view_moldinsight", "upload_file", "view_history", "view_inventory", "manage_inventory", "view_products", "manage_products", "view_suppliers", "manage_suppliers", "view_customers", "manage_customers", "view_finance", "manage_receipts", "manage_payments", "void_finance_transaction", "view_users", "manage_users", "manage_roles"]},
{"code": "user", "name": "普通用户", "description": "普通用户,可使用模具分析和查看库存", "is_system": False, "permissions": ["view_dashboard", "view_moldinsight", "upload_file", "view_history", "view_inventory", "view_products", "view_suppliers", "view_customers", "view_finance", "manage_receipts", "manage_payments"]},
{"code": "viewer", "name": "只读用户", "description": "只读用户,只能查看数据", "is_system": False, "permissions": ["view_dashboard", "view_moldinsight", "view_history", "view_inventory", "view_products", "view_suppliers", "view_customers", "view_finance"]},
]
async def init_permissions(session):
"""初始化权限"""
result = await session.execute(select(Permission))
existing_perms = result.scalars().all()
if existing_perms:
logger.info("权限已初始化")
return
perm_map = {}
for perm_data in DEFAULT_PERMISSIONS:
perm = Permission(**perm_data)
session.add(perm)
await session.flush()
perm_map[perm.code] = perm.id
logger.info(f"创建了 {len(DEFAULT_PERMISSIONS)} 个权限")
return perm_map
async def init_roles(session, perm_map):
"""初始化角色"""
result = await session.execute(select(Role))
existing_roles = result.scalars().all()
if existing_roles:
logger.info("角色已初始化")
return
for role_data in DEFAULT_ROLES:
perm_ids = [perm_map[code] for code in role_data.pop("permissions")]
role = Role(**role_data)
session.add(role)
await session.flush()
for perm_id in perm_ids:
rp = RolePermission(role_id=role.id, permission_id=perm_id)
session.add(rp)
logger.info(f"创建了 {len(DEFAULT_ROLES)} 个角色")
async def create_admin_user(session):
"""创建默认管理员"""
result = await session.execute(select(User).where(User.username == settings.ADMIN_USERNAME))
existing_admin = result.scalar_one_or_none()
if existing_admin:
logger.info("管理员账户已存在")
return
admin = User(
username=settings.ADMIN_USERNAME,
email=settings.ADMIN_EMAIL,
hashed_password=get_password_hash(settings.ADMIN_PASSWORD),
full_name=settings.ADMIN_FULL_NAME,
is_active=True
)
session.add(admin)
await session.flush()
result = await session.execute(select(Role).where(Role.code == "admin"))
admin_role = result.scalar_one_or_none()
if admin_role:
user_role = UserRole(user_id=admin.id, role_id=admin_role.id)
session.add(user_role)
await session.commit()
logger.info(f"创建了管理员账户: {settings.ADMIN_USERNAME}")
async def init_database(keep_connected: bool = True):
"""初始化数据库"""
try:
await db_manager.connect()
await db_manager.create_tables()
await ensure_schema_updates()
async with db_manager.session() as session:
perm_map = await init_permissions(session)
if perm_map is None:
# Permissions already existed, fetch them from database
result = await session.execute(select(Permission))
perms = result.scalars().all()
perm_map = {perm.code: perm.id for perm in perms}
await init_roles(session, perm_map)
await create_admin_user(session)
logger.info("数据库初始化完成")
print("=" * 60)
print("数据库初始化成功!")
print("=" * 60)
print(f"管理员用户名: {settings.ADMIN_USERNAME}")
print(f"管理员邮箱: {settings.ADMIN_EMAIL}")
print("=" * 60)
print("可以在 .env 文件中修改管理员配置:")
print(" ADMIN_USERNAME")
print(" ADMIN_EMAIL")
print(" ADMIN_FULL_NAME")
print("=" * 60)
return True
except Exception as e:
logger.error(f"数据库初始化失败: {e}")
print(f"数据库初始化失败: {e}")
return False
finally:
if not keep_connected:
await db_manager.disconnect()
async def ensure_schema_updates():
async with db_manager.engine.begin() as conn:
await conn.execute(text("ALTER TABLE products ADD COLUMN IF NOT EXISTS item_type VARCHAR(20) DEFAULT 'finished'"))
await conn.execute(text("UPDATE products SET item_type = 'finished' WHERE item_type IS NULL"))
await conn.execute(text("ALTER TABLE sales_orders ADD COLUMN IF NOT EXISTS production_status VARCHAR(20) DEFAULT 'not_started'"))
await conn.execute(text("ALTER TABLE sales_orders ADD COLUMN IF NOT EXISTS production_no VARCHAR(50)"))
await conn.execute(text("ALTER TABLE sales_orders ADD COLUMN IF NOT EXISTS planned_material_cost DOUBLE PRECISION DEFAULT 0"))
await conn.execute(text("ALTER TABLE sales_orders ADD COLUMN IF NOT EXISTS actual_material_cost DOUBLE PRECISION DEFAULT 0"))
await conn.execute(text("ALTER TABLE purchase_orders ADD COLUMN IF NOT EXISTS received_date TIMESTAMP WITHOUT TIME ZONE"))
await conn.execute(text("ALTER TABLE purchase_orders ADD COLUMN IF NOT EXISTS paid_date TIMESTAMP WITHOUT TIME ZONE"))
await conn.execute(text("""
CREATE TABLE IF NOT EXISTS product_materials (
id SERIAL PRIMARY KEY,
finished_product_id INTEGER NOT NULL REFERENCES products(id),
material_product_id INTEGER NOT NULL REFERENCES products(id),
quantity DOUBLE PRECISION NOT NULL,
loss_rate DOUBLE PRECISION DEFAULT 0,
created_at TIMESTAMP DEFAULT NOW(),
updated_at TIMESTAMP DEFAULT NOW()
)
"""))
await conn.execute(text("""
CREATE UNIQUE INDEX IF NOT EXISTS uq_product_material_unique
ON product_materials (finished_product_id, material_product_id)
"""))
await conn.execute(text("ALTER TABLE mold_cavity_data ADD COLUMN IF NOT EXISTS best_scheme_id VARCHAR(64)"))
await conn.execute(text("ALTER TABLE mold_cavity_data ADD COLUMN IF NOT EXISTS confidence_score DOUBLE PRECISION"))
await conn.execute(text("ALTER TABLE mold_cavity_data ADD COLUMN IF NOT EXISTS is_fallback BOOLEAN"))
await conn.execute(text("ALTER TABLE mold_cavity_data ADD COLUMN IF NOT EXISTS fallback_reason TEXT"))
await conn.execute(text("""
CREATE INDEX IF NOT EXISTS idx_mold_cavity_best_scheme_id
ON mold_cavity_data (best_scheme_id)
"""))
await conn.execute(text("""
CREATE INDEX IF NOT EXISTS idx_mold_cavity_is_fallback
ON mold_cavity_data (is_fallback)
"""))
await conn.execute(text("""
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'uq_inventory_product_warehouse'
) THEN
ALTER TABLE inventory
ADD CONSTRAINT uq_inventory_product_warehouse UNIQUE (product_id, warehouse_id);
END IF;
END $$;
"""))
await conn.execute(text("""
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ck_inventory_qty_nonnegative'
) THEN
ALTER TABLE inventory
ADD CONSTRAINT ck_inventory_qty_nonnegative
CHECK (quantity >= 0 AND locked_quantity >= 0 AND locked_quantity <= quantity);
END IF;
END $$;
"""))
await conn.execute(text("""
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ck_purchase_order_items_qty'
) THEN
ALTER TABLE purchase_order_items
ADD CONSTRAINT ck_purchase_order_items_qty
CHECK (quantity > 0 AND received_quantity >= 0 AND received_quantity <= quantity);
END IF;
END $$;
"""))
await conn.execute(text("""
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ck_sales_order_items_qty'
) THEN
ALTER TABLE sales_order_items
ADD CONSTRAINT ck_sales_order_items_qty
CHECK (quantity > 0 AND delivered_quantity >= 0 AND delivered_quantity <= quantity);
END IF;
END $$;
"""))
await conn.execute(text("ALTER TABLE sales_orders ADD COLUMN IF NOT EXISTS manufacturing_date TIMESTAMP WITHOUT TIME ZONE"))
await conn.execute(text("ALTER TABLE sales_orders ALTER COLUMN manufacturing_date TYPE TIMESTAMP WITHOUT TIME ZONE"))
if __name__ == "__main__":
asyncio.run(init_database(keep_connected=False))