# encoding: utf-8
__author__ = 'zhanghe'
from psycopg2 import *
import json
from datetime import date, datetime
# 数据库日志专用配置
from log import Logger
my_logger = Logger('postgres', 'postgres.log', 'DEBUG')
my_logger.set_file_level('DEBUG')
my_logger.set_stream_level('WARNING') # WARNING DEBUG
my_logger.set_stream_handler_fmt('%(message)s')
my_logger.load()
logger = my_logger.logger
# my_logger.get_memory_usage()
class Postgres(object):
"""
自定义Postgres工具
"""
def __init__(self, db_config, db_name=None):
self.db_config = db_config
if db_name is not None:
self.db_config['database'] = db_name
try:
self.conn = connect(
database=self.db_config['database'],
user=self.db_config['user'],
password=self.db_config['password'],
host=self.db_config['host'],
port=self.db_config['port']
)
except Exception, e:
logger.error(e)
@staticmethod
def __default(obj):
"""
支持datetime的json encode
TypeError: datetime.datetime(2015, 10, 21, 8, 42, 54) is not JSON serializable
:param obj:
:return:
"""
if isinstance(obj, datetime):
return obj.strftime('%Y-%m-%d %H:%M:%S')
elif isinstance(obj, date):
return obj.strftime('%Y-%m-%d')
else:
raise TypeError('%r is not JSON serializable' % obj)
def is_conn_open(self):
"""
检测连接是否打开
:return:
"""
if self.conn is None or self.conn.closed == 1:
return False
else:
return True
def close_conn(self):
"""
关闭数据库连接
:return:
"""
if self.is_conn_open() is True:
self.conn.close()
def truncate(self, table_name):
"""
清空表
:param table_name:
:return:
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return []
# 参数判断
if table_name is None:
logger.error('查询表名缺少参数')
return []
sql = 'truncate table %s' % table_name
logger.info(sql)
cursor = self.conn.cursor()
try:
cursor.execute(sql)
self.conn.commit()
logger.info('更新行数：%s' % cursor.rowcount)
cursor.close()
return True
except Exception, e:
logger.error(e)
finally:
cursor.close()
def get_columns_name(self, table_name):
"""
获取数据表的字段名称
:param table_name:
:return:
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return []
# 参数判断
if table_name is None:
logger.error('查询表名缺少参数')
return []
cursor = self.conn.cursor()
sql = "select column_name from information_schema.columns where table_name = '%s'" % table_name
logger.info(sql)
try:
cursor.execute(sql)
result = cursor.fetchall()
row = [item[0] for item in result]
return row
except Exception, e:
logger.error(e)
finally:
cursor.close()
def get_row(self, table_name, condition=None):
"""
获取单行数据
:return:
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return None
# 参数判断
if table_name is None:
logger.error('查询表名缺少参数')
return None
if condition and not isinstance(condition, list):
logger.error('查询条件参数格式错误')
return None
# 组装查询条件
if condition:
sql_condition = 'where '
sql_condition += ' and '.join(condition)
else:
sql_condition = ''
sql = 'select * from %s %s limit 1' % (table_name, sql_condition)
logger.info(sql)
cursor = self.conn.cursor()
try:
cursor.execute(sql)
row = cursor.fetchone()
return row
except Exception, e:
logger.error(e)
finally:
cursor.close()
def get_rows(self, table_name, condition=None, limit='limit 10 offset 0'):
"""
获取多行数据
con_obj.get_rows('company', ["type='6'"], 'limit 10 offset 0')
con_obj.get_rows('company', ["type='6'"], 'limit 10')
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return None
# 参数判断
if table_name is None:
logger.error('查询表名缺少参数')
return None
if condition and not isinstance(condition, list):
logger.error('查询条件参数格式错误')
return None
# 组装查询条件
if condition:
sql_condition = 'where '
sql_condition += ' and '.join(condition)
else:
sql_condition = ''
sql = 'select * from %s %s %s' % (table_name, sql_condition, limit)
logger.info(sql)
cursor = self.conn.cursor()
try:
cursor.execute(sql)
rows = cursor.fetchall()
return rows
except Exception, e:
logger.error(e)
finally:
cursor.close()
def get_count(self, table_name, condition=None):
"""
获取记录总数
:return:
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return 0
# 参数判断
if table_name is None:
logger.error('查询表名缺少参数')
return 0
if condition and not isinstance(condition, list):
logger.error('查询条件参数格式错误')
return 0
# 组装查询条件
if condition:
sql_condition = 'where '
sql_condition += ' and '.join(condition)
else:
sql_condition = ''
sql = 'select count(*) from %s %s' % (table_name, sql_condition)
logger.info(sql)
cursor = self.conn.cursor()
try:
cursor.execute(sql)
row = cursor.fetchone()
count = row[0]
return count
except Exception, e:
logger.error(e)
finally:
cursor.close()
def output_row(self, table_name, condition=None, style=0):
"""
格式化输出单个记录
style=0 键值对齐风格
style=1 JSON缩进风格
"""
# 参数判断
if not table_name:
logger.error('查询数据缺少参数')
return None
if condition and not isinstance(condition, list):
logger.error('查询条件参数格式错误')
return None
columns_name = self.get_columns_name(table_name)
row = self.get_row(table_name, condition)
if not columns_name:
logger.error('表名不存在')
return None
if not row:
logger.error('记录不存在')
return None
if style == 0:
# 获取字段名称最大的长度值作为缩进依据
max_len_column = max([len(each_column) for each_column in columns_name])
str_format = '{0: >%s}' % max_len_column
columns_name = [str_format.format(each_column) for each_column in columns_name]
result = dict(zip(columns_name, row))
print '********** 表名[%s] **********' % table_name
for key, item in result.items():
print key, ':', item
else:
result = dict(zip(columns_name, row))
print json.dumps(result, indent=4, ensure_ascii=False, default=self.__default)
def output_rows(self, table_name, condition=None, limit='limit 10 offset 0', style=0):
"""
格式化输出批量记录
style=0 键值对齐风格
style=1 JSON缩进风格
"""
# 参数判断
if not table_name:
logger.error('查询数据缺少参数')
return None
if condition and not isinstance(condition, list):
logger.error('查询条件参数格式错误')
return None
columns_name = self.get_columns_name(table_name)
rows = self.get_rows(table_name, condition, limit)
if not columns_name:
logger.error('表名不存在')
return None
if not rows:
logger.error('记录不存在')
return None
if style == 0:
# 获取字段名称最大的长度值作为缩进依据
max_len_column = max([len(each_column) for each_column in columns_name])
str_format = '{0: >%s}' % max_len_column
columns_name = [str_format.format(each_column) for each_column in columns_name]
count = 0
total = len(rows)
for row in rows:
result = dict(zip(columns_name, row))
count += 1
print '********** 表名[%s] [%d/%d] **********' % (table_name, count, total)
for key, item in result.items():
print key, ':', item
else:
for row in rows:
result = dict(zip(columns_name, row))
print json.dumps(result, indent=4, ensure_ascii=False, default=self.__default)
def update(self, table_name, update_field, condition=None):
"""
更新数据
con_obj.update('company', ["title='标题'", "flag='2'"], ["type='6'"])
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return False
# 参数判断
if not table_name or not update_field:
logger.error('更新数据缺少参数')
return False
if not isinstance(update_field, list) or (condition and not isinstance(condition, list)):
logger.error('更新数据参数格式错误')
return False
# 组装更新字段
if update_field:
sql_update_field = 'set '
sql_update_field += ' and '.join(update_field)
else:
sql_update_field = ''
# 组装更新条件
if condition:
sql_condition = 'where '
sql_condition += ' and '.join(condition)
else:
sql_condition = ''
# 拼接sql语句
sql = 'update %s %s %s' % (table_name, sql_update_field, sql_condition)
logger.info(sql)
cursor = self.conn.cursor()
try:
cursor.execute(sql)
self.conn.commit()
logger.info('更新行数：%s' % cursor.rowcount)
return True
except Exception, e:
logger.error(e)
finally:
cursor.close()
def delete(self, table_name, condition=None):
"""
删除数据
con_obj.delete('company', ["type='6'", "flag='2'"])
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return False
# 参数判断
if condition and not isinstance(condition, list):
logger.error('删除数据参数格式错误')
return False
# 组装删除条件
if condition:
sql_condition = 'where '
sql_condition += ' and '.join(condition)
else:
sql_condition = ''
# 拼接sql语句
sql = 'delete from %s %s' % (table_name, sql_condition)
logger.info(sql)
cursor = self.conn.cursor()
try:
cursor.execute(sql)
self.conn.commit()
logger.info('删除行数：%s' % cursor.rowcount)
logger.info('删除成功')
return True
except Exception, e:
logger.error(e)
finally:
cursor.close()
def query_by_sql(self, sql=None):
"""
根据sql语句查询
:return:
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return None
if sql is None:
logger.error('sql语句不能为空')
return None
# 安全性校验
sql = sql.lower()
logger.info(sql)
if not sql.startswith('select'):
logger.error('未授权的操作')
return None
cursor = self.conn.cursor()
try:
cursor.execute(sql)
rows = cursor.fetchall()
return rows
except Exception, e:
logger.error(e)
finally:
cursor.close()
def update_by_sql(self, sql=None):
"""
根据sql语句[增删改]
:return:
"""
if self.is_conn_open() is False:
logger.error('连接已断开')
return False
if sql is None:
logger.error('sql语句不能为空')
return False
# 安全性校验
sql = sql.lower()
logger.info(sql)
if not (sql.startswith('update') or sql.startswith('insert') or sql.startswith('delete')):
logger.error('未授权的操作')
return False
cursor = self.conn.cursor()
try:
cursor.execute(sql)
self.conn.commit()
logger.info('影响行数：%s' % cursor.rowcount)
logger.info('执行成功')
return True
except Exception, e:
logger.error(e)
finally:
cursor.close()