Dish_Recognition/dao/database.py

# coding=utf-8
"""双库架构：SQLite 本地临时缓存 + MySQL 总服务器。"""
import os
from contextlib import contextmanager

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker

from dao.db import Base
from dao.settings import CAIHUI_DATABASE, DISH_SQLITE_PATH

_local_engine = None
_local_session_factory = None
_server_engine = None
_server_session_factory = None
_server_available = False


def _mysql_uri():
    return (
        f"mysql+pymysql://{CAIHUI_DATABASE['user']}:{CAIHUI_DATABASE['password']}"
        f"@{CAIHUI_DATABASE['host']}:{CAIHUI_DATABASE['port']}/{CAIHUI_DATABASE['name']}"
        f"?charset=UTF8MB4"
    )


def _sqlite_uri():
    os.makedirs(os.path.dirname(DISH_SQLITE_PATH) or '.', exist_ok=True)
    return f"sqlite:///{DISH_SQLITE_PATH.replace(os.sep, '/')}"


def _import_models():
    import models  # noqa: F401


def _ensure_schema_columns(engine):
    """兼容旧库：补列、从 image_path 迁移到 image_url；bases 对齐管理系统字段。"""
    inspector = inspect(engine)
    tables = inspector.get_table_names()

    with engine.begin() as conn:
        if 'global_templates' in tables:
            existing = {col['name'] for col in inspector.get_columns('global_templates')}
            if 'image_url' not in existing:
                conn.execute(text('ALTER TABLE global_templates ADD COLUMN image_url VARCHAR(512)'))
            if 'thumbnail_url' not in existing:
                conn.execute(text('ALTER TABLE global_templates ADD COLUMN thumbnail_url VARCHAR(512)'))

            # 旧列数据迁到新列（若存在）
            existing = {col['name'] for col in inspect(engine).get_columns('global_templates')}
            if 'image_path' in existing:
                conn.execute(text(
                    'UPDATE global_templates SET image_url = image_path '
                    'WHERE (image_url IS NULL OR image_url = "") AND image_path IS NOT NULL'
                ))
            if 'thumbnail_path' in existing:
                conn.execute(text(
                    'UPDATE global_templates SET thumbnail_url = thumbnail_path '
                    'WHERE (thumbnail_url IS NULL OR thumbnail_url = "") AND thumbnail_path IS NOT NULL'
                ))

        if 'bases' in tables:
            _ensure_bases_columns(conn, engine)


def _ensure_bases_columns(conn, engine):
    """为已有 bases 表补齐管理系统同步字段（SQLite / MySQL 通用 ADD COLUMN）。"""
    existing = {col['name'] for col in inspect(engine).get_columns('bases')}
    dialect = engine.dialect.name
    adds = []
    if 'external_id' not in existing:
        adds.append(('external_id', 'INTEGER'))
    if 'branch_name' not in existing:
        adds.append(('branch_name', 'VARCHAR(100)'))
    if 'province' not in existing:
        adds.append(('province', 'VARCHAR(50)'))
    if 'city' not in existing:
        adds.append(('city', 'VARCHAR(50)'))
    if 'district' not in existing:
        adds.append(('district', 'VARCHAR(50)'))
    if 'address' not in existing:
        adds.append(('address', 'VARCHAR(255)'))
    if 'phone' not in existing:
        adds.append(('phone', 'VARCHAR(50)'))
    if 'description' not in existing:
        adds.append(('description', 'VARCHAR(255)'))
    if 'longitude' not in existing:
        adds.append(('longitude', 'FLOAT' if dialect == 'sqlite' else 'DOUBLE'))
    if 'latitude' not in existing:
        adds.append(('latitude', 'FLOAT' if dialect == 'sqlite' else 'DOUBLE'))
    if 'sort_order' not in existing:
        adds.append(('sort_order', 'INTEGER DEFAULT 0'))

    for name, coltype in adds:
        conn.execute(text(f'ALTER TABLE bases ADD COLUMN {name} {coltype}'))

    # base_code 可能原先 VARCHAR(50)，管理系统 code 较长；MySQL 可尝试加长
    if dialect == 'mysql' and 'base_code' in existing:
        try:
            conn.execute(text('ALTER TABLE bases MODIFY COLUMN base_code VARCHAR(128) NOT NULL'))
        except Exception:
            pass


def init_local_engine():
    global _local_engine, _local_session_factory
    if _local_engine is not None:
        return _local_engine
    _local_engine = create_engine(_sqlite_uri(), echo=False)
    _import_models()
    Base.metadata.create_all(_local_engine)
    _ensure_schema_columns(_local_engine)
    _local_session_factory = sessionmaker(bind=_local_engine, expire_on_commit=False)
    print("[OK] local cache sqlite connected")
    return _local_engine


def init_server_engine():
    global _server_engine, _server_session_factory, _server_available
    if _server_engine is not None:
        return _server_engine
    try:
        engine = create_engine(
            _mysql_uri(),
            echo=False,
            pool_recycle=3600,
            pool_pre_ping=True,
        )
        with engine.connect() as conn:
            conn.execute(text('SELECT 1'))
        _import_models()
        Base.metadata.create_all(engine)
        _ensure_schema_columns(engine)
        _server_engine = engine
        _server_session_factory = sessionmaker(bind=_server_engine, expire_on_commit=False)
        _server_available = True
        print("[OK] server mysql connected")
        return _server_engine
    except Exception as exc:
        _server_available = False
        print(f"[WARN] server mysql unavailable: {exc}")
        return None


def is_server_available():
    if _server_engine is None:
        init_server_engine()
    return _server_available


@contextmanager
def local_session_scope():
    if _local_session_factory is None:
        init_local_engine()
    session = _local_session_factory()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def server_session_scope():
    if not is_server_available():
        raise RuntimeError("MySQL 总服务器不可用")
    session = _server_session_factory()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def get_backend():
    return 'sqlite(cache)+mysql(server)' if is_server_available() else 'sqlite(cache-only)'