python上下文小记

python访问数据库是日常编程中经常使用的,一般的代码如下:

可以看出样例代码还是比较多的,有没有办法优化一下呢?有!

def send_msgs(conn_id=None, **kwargs):
    conn = get_mysql_conn(conn_id)
    cursor = conn.cursor()
    sql_db = '''select count(1) from table)'''
    cursor.execute(sql_db)
    result = cursor.fetchone()
    cursor.close()
    conn.close()

采用python提过的上下文方式,访问数据库就可以这样写:

def send_msgs(conn_id=None, **kwargs):
    with db_conn(conn) as cur:
        cur.execute(sql)
        r1=cur.fetchone()

缩减了一半以上的代码,而且还可以自动关闭,db_conn怎么写?就是实现下面这样一个类就可以了:

class db_conn(object):

    def __init__(self, conn_id):
        self.conn = MySqlHook.get_connection(conn_id=conn_id).get_hook().get_conn()
        self.cur = self.conn.cursor()

    def __enter__(self):
        return self.cur

    def __exit__(self, exc_type, exc_value, exc_trace):
        self.conn.commit()
        self.cur.close()
        self.conn.close()