在python中给“数据库连接”，请个保姆

Author : Winter Han

对于新手来说，数据库连接的管理，是个险活，一旦漏关，会导致数据库服务器“连接过多”。
对于老手来说，又是个累活，每个数据库操作的function中，同样的代码都要写一遍。

怎么破？
我们尝试用 with （上下文管理器）来解决。

先说说with(上下文管理器)

可以在with的body，的前后做一些业务操作。有点像AOP（面向方面编程）思想，或者 python中的装饰器。
with语句返回的对象，要包括 “__enter__”, "__exit__"两个方法，分别是在 body之前执行，body之后执行。

先看个示例：

# -*- coding: utf-8 -*-

class PrintContext(object):  

    def __enter__(self):
        print "info: at __enter__"

    def __exit__(self, exc, value, tb):
        print "info: at __exit__"


if __name__ == "__main__":
    with PrintContext() as p:
        print "info: at body"

输出：

info: at __enter__
info: at body
info: at __exit__

传统的数据库连接方式

# -*- coding: utf-8 -*-
import MySQLdb


def get_all_user():
    conn = MySQLdb.connect(HOST, USERNAME, PASSWORD, DATABASE)
    cursor = conn.cursor()

    sql = "select * from user"
    cursor.execute(sql)

    result = cursor.fetchall()

    cursor.close()
    conn.close()

    return result

大家可以看到，其实真正的业务逻辑，就中间三行而已。

用with改进后的数据库连接

# -*- coding: utf-8 -*-
import MySQLdb

HOST, USERNAME, PASSWORD, DATABASE = ("HOST", "USERNAME", "PASSWORD", "DATABASE")

class DBCloser(object):
    """docstring for DBCloser"""

    def __init__(self, connect, cursorclass=None):
        self.connect = connect
        self.cursor = None
        self.cursorclass = cursorclass

    def __enter__(self):
        self.cursor = self.connect.cursor(cursorclass=self.cursorclass)
        return self.cursor

    def __exit__(self, exc, value, tb):
        if exc:
            self.connect.rollback()
        else:
            self.connect.commit()

        self.cursor.close()
        self.connect.close()


def get_all_user():
    conn = DBCloser(MySQLdb.connect(HOST, USERNAME, PASSWORD, DATABASE))
    
    with conn as cursor:
        sql = "select * from user"
        cursor.execute(sql)
        result = cursor.fetchall()    

    return result

这样的好处是：

一、为每个function 节省 两行代码；
二、支持rollback
三、提高代码扩展性。共性的逻辑，可以加在__enter__, __exit__ 之中。