QingLong/AI/Text2Sql/Model/biModel.py

from pydantic import BaseModel
from sqlalchemy import create_engine, Column, Integer, String, SmallInteger, Date
from sqlalchemy.orm import declarative_base  # 更新导入路径
from sqlalchemy.orm import sessionmaker

# 数据库连接配置
DATABASE_URL = "postgresql+psycopg2://postgres:DsideaL147258369@10.10.14.71:5432/szjz_db"

# 创建数据库引擎
engine = create_engine(DATABASE_URL)

# 创建 SessionLocal 类
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建 Base 类
Base = declarative_base()

# 获取数据库会话
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 定义 t_bi_question 表模型
class TBIQuestion(Base):
    __tablename__ = "t_bi_question"

    id = Column(Integer, primary_key=True, index=True)
    question = Column(String(255), nullable=False)
    state_id = Column(SmallInteger, nullable=False, default=0)
    sql = Column(String(2048))
    bar_chart_x_columns = Column(String(255))
    bar_chart_y_columns = Column(String(255))
    pie_chart_category_columns = Column(String(255))
    pie_chart_value_column = Column(String(255))
    session_id = Column(String(255), nullable=False)
    excel_file_name = Column(String(255))
    bar_chart_file_name = Column(String(255))
    pie_chart_file_name = Column(String(255))
    report_file_name = Column(String(255))
    create_time = Column(Date, nullable=False, default="CURRENT_TIMESTAMP")
    is_system = Column(SmallInteger, nullable=False, default=0)
    is_collect = Column(SmallInteger, nullable=False, default=0)

# 创建 Pydantic 模型
class TBIQuestionCreate(BaseModel):
    question: str
    state_id: int = 0
    sql: str = None
    bar_chart_x_columns: str = None
    bar_chart_y_columns: str = None
    pie_chart_category_columns: str = None
    pie_chart_value_column: str = None
    session_id: str
    excel_file_name: str = None
    bar_chart_file_name: str = None
    pie_chart_file_name: str = None
    report_file_name: str = None
    is_system: int = 0
    is_collect: int = 0

class TBIQuestionUpdate(BaseModel):
    question: str = None
    state_id: int = None
    sql: str = None
    bar_chart_x_columns: str = None
    bar_chart_y_columns: str = None
    pie_chart_category_columns: str = None
    pie_chart_value_column: str = None
    session_id: str = None
    excel_file_name: str = None
    bar_chart_file_name: str = None
    pie_chart_file_name: str = None
    report_file_name: str = None
    is_system: int = None
    is_collect: int = None