# coding:utf-8
from tools_hjh.ThreadPool import ThreadPool


class DBConn:
    """ 维护一个关系型数据库连接池，目前支持oracle，pgsql，mysql，sqlite；支持简单的sql执行 """
    
    ORACLE = 'oracle'
    PGSQL = 'pgsql'
    MYSQL = 'mysql'
    SQLITE = 'sqlite'

    def __init__(self, dbtype, host=None, port=None, db=None, username=None, password=None, poolsize=2, encoding='UTF-8', lib_dir=None):
        """ 初始化连接池
                如果是sqlite，db这个参数是要显示给入的，内存数据库 db=':memory:'
                如果是oracle，db给入的是sid或是servername都是可以的 """
                
        self.dbtype = dbtype
        self.host = host
        self.port = port
        self.db = db
        self.username = username
        self.password = password
        self.poolsize = poolsize
        self.encoding = encoding
        self.lib_dir = lib_dir
        
        self.runtp = ThreadPool(1, save_result=True)
        self.dbpool = None
        
        self.config = {
            'host':self.host,
            'port':self.port,
            'database':self.db,
            'user':self.username,
            'password':self.password,
            'maxconnections':self.poolsize,  # 最大连接数
            'blocking':True,  # 连接数达到最大时，新连接是否可阻塞
            'reset':False
        }
        
        if self.dbtype == 'pgsql' or self.dbtype == 'mysql':
            from dbutils.pooled_db import PooledDB
        if self.dbtype == "pgsql":
            import psycopg2
            self.dbpool = PooledDB(psycopg2, **self.config)
        elif self.dbtype == "mysql":
            import pymysql
            self.dbpool = PooledDB(pymysql, **self.config)
        elif self.dbtype == "sqlite":
            if self.db is None:
                self.db = ':memory:'
            import sqlite3
            from dbutils.persistent_db import PersistentDB
            self.dbpool = PersistentDB(sqlite3, database=self.db)
        elif self.dbtype == "oracle":
            import cx_Oracle
            if lib_dir is not None:
                cx_Oracle.init_oracle_client(lib_dir=lib_dir)
            try:
                dsn = cx_Oracle.makedsn(host, port, service_name=self.db)
                self.dbpool = cx_Oracle.SessionPool(user=username,
                                                    password=password,
                                                    dsn=dsn,
                                                    max=poolsize,
                                                    increment=1,
                                                    encoding=encoding)
            except:
                dsn = cx_Oracle.makedsn(host, port, sid=self.db)
                self.dbpool = cx_Oracle.SessionPool(user=username,
                                                    password=password,
                                                    dsn=dsn,
                                                    max=poolsize,
                                                    increment=1,
                                                    encoding=encoding)
    
    def run_procedure(self, sql, param=None, auto_commit=None):
        # 替换占位符
        if self.dbtype == 'pgsql' or self.dbtype == 'mysql':
            sql = sql.replace('?', '%s')
        elif self.dbtype == 'oracle':
            sql = sql.replace('?', ':1')            
        else:
            pass
                
        # 获取连接
        if self.dbtype == "oracle":
            conn = self.dbpool.acquire()
        else:
            conn = self.dbpool.connection()
            
        cur = conn.cursor()
                
        # 执行非SELECT语句
        sql = sql.strip()
        if type(param) == list:
            cur.executemany(sql, param)
        elif type(param) == tuple:
            cur.execute(sql, param)
        elif param is None:
            cur.execute(sql)
        if auto_commit: 
            if sql.lower().strip().startswith("update") or sql.lower().strip().startswith("delete") or sql.lower().strip().startswith("insert"):
                conn.commit()
        rownum = cur.rowcount
        rs = rownum
        cur.close()
        conn.close()
        return rs
    
    def __run(self, sqls, param=None, auto_commit=None):

        # 替换占位符
        if self.dbtype == 'pgsql' or self.dbtype == 'mysql':
            sqls = sqls.replace('?', '%s')
        elif self.dbtype == 'oracle':
            sqls = sqls.replace('?', ':1')            
        else:
            pass
        
        sql_list = []
        
        # sql只有一行
        if not '\n' in sqls:
            sql_list.append(sqls.rstrip(';'))
            
        # sql有多行
        else:
            # 去掉每行的首尾空格、换行，再去掉最后一个;,去掉--开头的行
            str2 = ''
            for line in sqls.split('\n'):
                line = line.strip()
                if not line.startswith('--') and len(line) > 0:
                    str2 = str2 + line + '\n'
            for sql in str2.split(';\n'):
                if sql is not None and sql != '' and len(sql) > 0:
                    sql_list.append(sql)
        
        # 获取连接
        if self.dbtype == "oracle":
            conn = self.dbpool.acquire()
        else:
            conn = self.dbpool.connection()
            
        cur = conn.cursor()
        
        for sql in sql_list:
            # 执行SELECT语句
            if sql.lower().strip().startswith("select") or (sql.lower().strip().startswith("with") and 'select' in sql.lower()):
                sql = sql.strip()
                if param is None:
                    cur.execute(sql)
                elif type(param) == tuple:
                    cur.execute(sql, param)
                rs = QueryResults(cur, conn)
                
            # 执行非SELECT语句
            elif not sql.lower().strip().startswith("select"):
                sql = sql.strip()
                if type(param) == list:
                    cur.executemany(sql, param)
                elif type(param) == tuple:
                    cur.execute(sql, param)
                elif param is None:
                    cur.execute(sql)
                if auto_commit: 
                    if sql.lower().strip().startswith("update") or sql.lower().strip().startswith("delete") or sql.lower().strip().startswith("insert"):
                        conn.commit()
                rownum = cur.rowcount
                rs = rownum
                cur.close()
                conn.close()
        # cur.close()
        # conn.close()
        return rs

    def run(self, sql, param=None, wait=False, auto_commit=True):
        """ 执行点什么
        sql中的占位符可以统一使用“?”
        wait为True则会等待当前正在执行的sql，有bug，暂不处理，自用规避"""
        if wait == True:
            tpnum = self.runtp.run(self.__run, (sql, param, auto_commit))
            self.runtp.wait()
            rs = self.runtp.result_map.pop(tpnum)
            return rs
        else:
            return self.__run(sql, param, auto_commit)
        
    def insert(self, table_name, rows, pk_cols=''):
        """ 往指定table_name中插入数据，rows是一个多个元组的列表，每个元组表示一组参数；或者是一个元组 """
        if type(rows) == list and len(rows) > 0:
            row = rows[0]
        elif type(rows) == tuple:
            row = rows
        elif len(rows) == 0:
            return 0
        
        sql2 = ''
        pk_num = 0
        if ',' not in pk_cols and pk_cols != '':
            pk_num = 1
            sql2 = pk_cols + ' = ?'
        elif ',' in pk_cols and len(pk_cols.split(',')) == 1:
            pk_num = 1
            sql2 = pk_cols.strip().strip(',') + ' = ?'
        elif ',' in pk_cols and len(pk_cols.split(',')) > 1:
            pk_num = len(pk_cols.split(','))
            for pk in pk_cols.split(','):
                sql2 = sql2 + pk + ' = ? and '
            sql2 = sql2.rstrip('and ')
            
        param_num = '?'
        for _ in range(len(row) - 1 - pk_num):
            param_num = param_num + ', ?'
            
        if self.dbtype == 'oracle' or self.dbtype == 'mysql':
            if len(pk_cols) == 0:
                sql = 'insert into ' + table_name + ' select ' + param_num + ' from dual'
            else:
                sql = 'insert into ' + table_name + ' select ' + param_num + ' from dual where not exists(select 1 from ' + table_name + ' where ' + sql2 + ')'
        else:
            if len(pk_cols) == 0:
                sql = 'insert into ' + table_name + ' select ' + param_num
            else:
                sql = 'insert into ' + table_name + ' select ' + param_num + ' where not exists(select 1 from ' + table_name + ' where ' + sql2 + ')'
        
        return self.run(sql, rows)
    
    def close(self):
        try:
            self.dbpool.close()
        except:
            pass
        finally:
            self.dbpool = None
    
    def __del__(self):
        self.close()
    
        
class QueryResults:

    def __init__(self, cur, conn):
        self.cur = cur
        self.conn = conn
        self.row = ''

    def get_cols(self):
        col = []
        for c in self.cur.description:
            col.append(c[0])
        return tuple(col)

    def get_rows(self, rownum='all'):
        if rownum == 'all' or rownum == True or rownum == False:
            rows = self.cur.fetchall()
        elif int(rownum) == 1:
            rows = self.cur.fetchone()
            if len(rows) == 0:
                self.close()
        else:
            rows = self.cur.fetchmany(int(rownum))
            if len(rows) == 0:
                self.close()
            
        rows_new = []
        for row in rows:
            row_new = []
            for cell in row:
                if 'LOB' not in str(type(cell)):
                    cell_new = cell
                else:
                    cell_new = cell.read()
                row_new.append(cell_new)
            rows_new.append(row_new)
            
        if rownum == 'all' or rownum == True or rownum == False:
            self.close()
        
        return rows_new
    
    def next(self):
        self.row = self.cur.fetchone()
        if self.row == None:
            self.close()
            return False
        else:
            return True

    def get_row(self):
        row_new = []
        for cell in self.row:
            if 'LOB' not in str(type(cell)):
                cell_new = cell
            else:
                cell_new = cell.read()
            row_new.append(cell_new)
        return row_new

    def close(self):
        try:
            self.cur.close()
        except:
            pass
        try:
            self.conn.close()
        except:
            pass
        
    def __del__(self):
        self.close()

