postgre与Cassandra数据库封装

目前公司使用的是Postgre和Cassandra数据库，python连接数据库写相关脚本，方便查询数据实现自动化脚本断言sql

#首先导入第三方包
import psycopg2.extras
import psycopg2.extensions

from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster

#这是pg库的登录参数，具体信息参考公司的
DB_CONFIG_PG =
{
'dbname':'device-service',
'user':'xx',
'password':'xx',
'host':'xx,
'port':1921
}


class Postgre:

    def __init__(self, db_config):
        """
        创建postgre库连接
        """
        try:
            self.conn = psycopg2.connect(**DB_CONFIG_PG)
        except Exception as e:

            raise e

    def get_one(self, sql):

        with self.conn:
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
                curs.execute(sql)
                return curs.fetchone()
        # cu = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        # cu.execute(sql)
        # return cu.fetchone()

    def get_many(self, sql, num):
        with self.conn:
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
                curs.execute(sql, num)
                return curs.fetchmany()

    def get_all(self, sql):
        """
        获取所有数据
        """
        with self.conn:
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
                curs.execute(sql)
                return curs.fetchall()

    def exist(self, sql):
        """
        判断是否存在
        """
        with self.conn:
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
                curs.execute(sql)
                if curs.fetchone():
                    return True
                else:
                    return False

    def delete(self, sql):
        with self.conn:
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
                curs.execute(sql)
        self.conn.commit()

    def __del__(self):
        """
        关闭连接
        """
        self.conn.close()
class Cassandra:

    def __init__(self):
        # 连接数据库
        try:
            node_ips = ["xx"]
            auth_provider = PlainTextAuthProvider(username='', password='')
            self.cluster = Cluster(node_ips, port=9042, auth_provider=auth_provider, connect_timeout=6000,)
            self.session = self.cluster.connect()

        except Exception as e:

            raise e

    #  查询数据
    def can_select(self, sql):
        rows = self.session.execute(sql)
        return rows

    #  新增，插入数据
    def can_insert(self, sql, data):
        self.session.execute(sql, data)

    # def find_count(self,sql):
    #     """
    #     查找返回结果的数量
    #     """
    #     res = self.session.execute(sql)
    #     return res

    def can_delete(self, sql):
        self.session.execute(sql)

    def __del__(self):
        """
        关闭游标，断开连接
        """
        self.session.shutdown()
        self.cluster.shutdown()

#调试
if __name__ == '__main__':
    import settings

    # db = Postgre(db_config=settings.DB_CONFIG_PG)
    # sql = 'SELECT t.*FROM public.edge_device t LIMIT 501'
    # res = db.get_one(sql)
    # print(res)
    a = Cassandra()
    res = a.can_select("select *from dbpe.point where key=142270 and ts>'2021-03-04 08:18:30.000' and week=267")
    print(res.one())