python 学习笔记（七）：mysql 数据库连接

不废话，直接贴代码

#!/usr/bin/python
# -*- coding: utf-8 -*-
from __future__ import print_function
import os
from contextlib import contextmanager

'''根据环境变量导入不同的数据库模块'''
if os.getenv('DB','MySQL') == 'mysql':
    import MySQLdb as db
else:
    import sqlite3 as db


@contextmanager
def get_conn(**kwargs):
    try:
        if os.getenv('DB','MySQL')=='mysql':
            conn = db.connect(
                host=kwargs.get('host', 'localhost'),
                user=kwargs.get('user'),
                passwd=kwargs.get('passwd'),
                port=kwargs.get('port', 3306),
                db=kwargs.get('db')
            )
        else:
            conn = db.connect(database=kwargs.get('db'))
        yield conn
    finally:
        if conn:
            conn.close()


def main():
    try:
        conn = dict(
            host='127.0.0.1',
            user='root',
            passwd='root',
            db='sd'
        )
        with get_conn(conn) as conn:
            with conn as cur:
                cur.execute('select * from ta_users;')

    except NotImplementedError as err:
        print(err)
    finally:
        cur.close()
        conn.close()