from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

class Database:
    def __init__(self, db_url):
        self.engine = create_engine(db_url)
        self.Session = sessionmaker(bind=self.engine)

    def query(self, sql_query):
        with self.Session() as session:
            result = session.execute(sql_query)
            rows = result.fetchall()

        return rows

    def insert(self, sql_query, **params):
        with self.Session() as session:
            session.execute(sql_query, params)
            session.commit()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.Session.close_all()
        self.engine.dispose()

在上述示例中,我们定义了一个 Database 类,用于执行数据库操作,并自动关闭数据库连接。该类包含了两个方法:query 和 insert,分别用于执行查询和插入操作。该类还实现了上下文管理器协议,可以在 with 语句中使用,并在退出上下文管理器时自动关闭数据库连接。

在类的 init 方法中,我们首先创建了一个数据库引擎对象,并使用 sessionmaker 类创建了一个会话工厂(SessionFactory)。在执行查询和插入操作时,我们创建了一个数据库会话对象,并在 with 语句中执行数据库操作。在退出 with 语句时,会自动调用会话对象的 close() 方法,从而自动关闭数据库连接。

下面是一个调用示例:

db = Database('mysql+pymysql://user:password@host:port/database')

# 查询数据
rows = db.query('SELECT * FROM mytable')
for row in rows:
    print(row)

# 插入数据
db.insert('INSERT INTO mytable (id, name) VALUES (:id, :name)', id=1, name='John')

# 自动关闭数据库连接
with db:
    # 执行数据库操作
    pass

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐