0. 前言

最近项目中需要用到kudu, 理论上最正确的方式是使用impala来查询一些聚合数据返回, 但是因为业务的一些性能要求, 如果使用impala 连接会将性能要求堵在impala上, 所以选择自己开发了一个kudu的连接池

1. 开发思路梳理

数据库的连接池有两个最大的问题, 连接失效和线程安全. 线程安全的问题可以用python queue模版中的Queue对象解决(https://docs.python.org/3.7/library/queue.html), 所以连接失效的处理就是这份连接池代码的核心和难点所在

连接失效的情况有两种, 一种是项目启动时连接失效, 另一种是使用的时候失效, 处理的思路如下:

  1. 项目启动时失效: 项目启动的时候会先创建一个连接, 如果创建失败, 则抛出异常, 整个项目停止.
  2. 项目运行中失效: 项目运行中检查失效的操作是从连接池中获取连接后先ping一下, 如果ping不通, 则新建一个连接, 进行操作后再放入连接池

根据这个处理的思路, 我们可以抽象新建连接和获取连接的的代码

创建代码的逻辑: 可以重试三次, 如果不成功返回None

def _create_new_conn(self):
    # 创建成功则返回链接, 否则返回None
    try_times = 0
    while try_times < 3:
        try:
            conn = kudu.connect(
                host=self.kwargs.get("host"), port=self.kwargs.get("port"))
            return conn
        except kudu.KuduBadStatus:
            logging.exception("kudu try connect error,please check kudu config or server")
            try_times += 1
            continue
    return None

获取连接的逻辑: 拿到连接池的一个连接后先ping一下, 如果ping不通则新建一个然后返回, 这里需要注意的是新建也有可能还是一个None, 这里处理的逻辑是使用的乐观处理, 即不理会这次生成的是什么, 直接将新建的对象返回

def _get_conn(self):
    conn = self.conn_queue.get()
    if conn is None:
        # 乐观处理, 如果为None就再新建一次
        return self._create_new_conn()

    # ping一下服务器, 看当前链接是否可用, 如果不可用也重新创建一次
    table = conn.list_tables()
    if not table:
        return self._create_new_conn()
    return conn

由于乐观处理的逻辑, 我们需要在创建连接池的时候先判断一次第一个创建的连接, 如果第一次就创建失败, 那么直接报错然后抛出异常, 因为一般这种情况大概率是配置错误相关的原因

first_conn = self._create_new_conn()
if not first_conn:
    raise Exception("Kudu first connect fail, please check kudu config or server")

其次, 在使用过程中则忽视连接, 直接进行操作, 然后进行对应的错误处理, 所以, 这个代码对日志的监控很重要, 出错的监控一定要认真观察, 防止数据库崩溃但是所有人都没感知.

2. kudu连接池完整源码

import queue
import logging
import kudu


class MyKudu:
    def __init__(self, **kwargs):
        self.size = kwargs.get('size', 10)
        self.kwargs = kwargs
        self.conn_queue = queue.Queue(maxsize=self.size)
        # 第一次链接如果失败则停止项目启动
        first_conn = self._create_new_conn()
        if not first_conn:
            raise Exception("Kudu first connect fail, please check kudu config or server")
        self.conn_queue.put(first_conn)
        for i in range(self.size-1):
            self.conn_queue.put(self._create_new_conn())

    async def close(self):
        try:
            while True:
                conn = self.conn_queue.get_nowait()
                if conn:
                    conn.close()
        except queue.Empty:
            pass
        logging.info("kudu connect quit")

    """
    kudu链接池代码
    """
    def _create_new_conn(self):
        # 创建成功则返回链接, 否则返回None
        try_times = 0
        while try_times < 3:
            try:
                conn = kudu.connect(
                    host=self.kwargs.get("host"), port=self.kwargs.get("port"))
                return conn
            except kudu.KuduBadStatus:
                logging.exception("kudu try connect error,please check kudu config or server")
                try_times += 1
                continue
        return None

    def _put_conn(self, conn):
        self.conn_queue.put(conn)

    def _get_conn(self):
        conn = self.conn_queue.get()
        if conn is None:
            # 乐观处理, 如果为None就再新建一次
            return self._create_new_conn()

        # ping一下服务器, 看当前链接是否可用, 如果不可用也重新创建一次
        table = conn.list_tables()
        if not table:
            return self._create_new_conn()
        return conn

    def _get_table(self, conn, table_name):
        try:
            table = conn.table(table_name)
        except kudu.errors.KuduNotFound:
            logging.exception("kudu table not found, please check kudu config or server")
            self._put_conn(conn)
            return None
        except AttributeError:
            logging.exception("kudu connect error, please check kudu config or server")
            self._put_conn(conn)
            return None
        return table

    def search(self, table_name, params: dict):
        conn = self._get_conn()
        table = self._get_table(conn, table_name)
        if table is None:
            return None

        scanner = table.scanner()
        for key in params:
            try:
                scanner.add_predicate(table[key] == params[key])
            except KeyError:
                logging.exception(f"kudu table key not found: {key}")
                self._put_conn(conn)
                return None

        result = scanner.open().read_all_tuples()
        self._put_conn(conn)
        return result

    def upsert(self, table_name, params: dict):
        conn = self._get_conn()
        table = self._get_table(conn, table_name)
        if table is None:
            return None

        session = conn.new_session()
        op = table.new_upsert(params)
        session.apply(op)
        try:
            session.flush()
        except kudu.KuduBadStatus:
            logging.exception(f"kudu upsert error, {table_name}: {params}")
        finally:
            self._put_conn(conn)