from utils.rdbmsUtil.databaseUtil import DatabaseUtil
from utils.rdbmsUtil.modelBase import Base
from utils.rdbmsUtil.modelUser import User, user_factory  # user_factory 确保工厂注册代码被执行
from utils.rdbmsUtil.modelUserOperations import UserOperations
# company_factory 确保工厂注册代码被执行
from utils.rdbmsUtil.modelCompany import Company, company_factory
from utils.rdbmsUtil.modelCompanyOperations import CompanyOperations

from sqlalchemy.sql import text
import pandas as pd


def access_sqlite():
    db = DatabaseUtil('sqlite:///mydatabase.db')
    Base.metadata.create_all(db.engine)

    companys = [Company(id=1, name="DS-Tech"), Company(id=2, name="ZN-Tech")]
    users = [User(name="Alice", age=31, companyid=1),
            User(name="Bob", age=25, companyid=2)]


    with db.session_scope() as session:
        company_ops = CompanyOperations()

        user_ops = UserOperations()
        
        company_ops.upsert_companys(session, companys)
        user_ops.upsert_users(session, users)
        # session.flush()  # 确保所有的 SQL 操作都被推送到数据库

        user_ops.update(session, name="Alice", age=32)

        result = user_ops.get_user(session, name="Alice", age=32)
        for user in result:
            print(user)


def access_tidb():
    print('access TiDB')
    DatabaseUtil(url='mysql+pymysql://root:123456@192.168.50.234:4000/energy',
                        pool_size=5,
                        max_overflow=3,
                        connect_timeout=1800)
    db = DatabaseUtil('mysql+pymysql://root:123456@192.168.50.234:4000/energy')

    with db.session_scope() as session:
        query_text = f"""SELECT * FROM analysis_result  """
        query_result = session.execute(text(query_text)).fetchall()

        dataFrame = pd.DataFrame(query_result)

        print(dataFrame.head(5))

access_tidb()