albin3 / book-notes Goto Github PK

import socks
import socket

# Can be socks4/5
socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS4,'127.0.0.1', 9050)
socket.socket = socks.socksocket

# Magic!
def getaddrinfo(*args):
    return [(socket.AF_INET, socket.SOCK_STREAM, 6, '', (args[0], args[1]))]

# do some requests

#!/usr/bin/env python
# -*- coding=utf-8 -*-
# @Author  : magong
# @Time    : 2022/5/31

import os
import logging

from impala.dbapi import connect
from krbcontext import krbcontext
import pandas as pd
from pydash import py_
from traffic_util.connector.base_query import IBaseQuery
from traffic_util.util import get_logger

from config import config

logger = get_logger(__name__)


class HiveClient(IBaseQuery):
    __instance = None

    @classmethod
    def instance(cls, *args, **kwargs):
        if cls.__instance:
            return cls.__instance
        else:
            try:
                with krbcontext(using_keytab=True, principal=py_.get(args[0], 'principal'), keytab_file=py_.get(args[0], 'keytab_file_path')):
                    logger.warning('Make sure `krb5.conf` is placed in `/etc` directory and host is added in `/etc/hosts`. ')
                    conn = connect(host=py_.get(args[0], 'host'),
                                   port=py_.get(args[0], 'port'),
                                   auth_mechanism=py_.get(args[0], 'auth_mechanism'),
                                   kerberos_service_name=py_.get(args[0], 'kerberos_service_name'),
                                   database=py_.get(args[0], 'database'))
                    return cls(conn)
            except Exception as e:
                # return None
                raise e

    def __init__(self, conn):
        self.conn = conn

    def execute(self, sql):
        """执行."""
        try:
            cur = self.conn.cursor()
            cur.execute('set hive.exec.dynamic.partition.mode=nonstrict')
            cur.execute('set hive.exec.max.dynamic.partitions.pernode=10000')
            cur.execute('set hive.exec.max.dynamic.partitions=10000')
            cur.execute(sql)
        except Exception as err:
            print(err)
            raise err

    def query(self, sql):
        """查询."""
        cur = self.conn.cursor()
        res = None
        try:
            cur.execute('set hive.exec.dynamic.partition.mode=nonstrict')
            cur.execute('set hive.exec.max.dynamic.partitions.pernode=10000')
            cur.execute('set hive.exec.max.dynamic.partitions=10000')
            cur.execute(sql)
            res = cur.fetchall()
        except Exception as err:
            print("查询失败, %s" % err)
            raise err
        finally:
            return res
    
    def execute_with_result(self, sql):
        logger.info('''
            executing sql: >>> {sql} <<<
            '''.format(sql=sql))
        
        cur = self.conn.cursor()
        cur.execute('set hive.exec.dynamic.partition.mode=nonstrict')
        cur.execute('set hive.exec.max.dynamic.partitions.pernode=10000')
        cur.execute('set hive.exec.max.dynamic.partitions=10000')
        cur.execute(sql)
        res = cur.fetchall()
        # cur.close()

        return cur, res

    def create_table(self):
        pass

    def exist_table(self):
        pass

    def query_to_df(self, sql):
        """查询输出DataFrame."""
        with self.conn.cursor() as cursor:
            logger.info('''
                executing sql: >>> {sql} <<<
                '''.format(sql=sql))
            
            cursor.execute('set hive.exec.dynamic.partition.mode=nonstrict')
            cursor.execute('set hive.exec.max.dynamic.partitions.pernode=10000')
            cursor.execute('set hive.exec.max.dynamic.partitions=10000')
            cursor.execute(sql)
            columns = [col[0] for col in cursor.description]
            records = [dict(zip(columns, row)) for row in cursor.fetchall()]
            results = pd.DataFrame(records)
            results.columns = columns
            return results
    
    def close(self):
        self.conn.close()
        
    def execute_sql(self, sql, sql_hints=None, show_log=False):
        if sql.endswith(';'):
            sql = sql[:-1]
        if show_log:
            logger.info('''
executing sql: >>> {sql} <<<
'''.format(sql=sql))
        return self.execute(sql)
    
    def run_sql_return_plain_json(self, sql):
        result = []

        cur, fetched_list = self.execute_with_result(sql)
        if cur.description is None or len(cur.description) < 1:
            return result

        columns = [col[0] for col in cur.description]
        for row in fetched_list:
            meta = {}
            for col, val in zip(columns, row):
                meta[col] = val
            result.append(meta)

        return result
    
    def run_sql_with_logview_return_plain_json(self, sql):
        pass 
    
    def check_table_exists(self, table_name):
        '''检查表是否存在.'''
        # TODO: 
        print(table_name)
        query_res = self.query('''SELECT COUNT(*) FROM information_schema.tables WHERE table_name = '{0}' '''.format(table_name.replace('\'', '\'\'')))
        if query_res[0] == 1:
            print(true) 
        # return len(self.query('''SHOW TABLES LIKE '{table_name}' '''.format(table_name=table_name))) == 1
    
    def truncate_table_if_exists(self, table_name, force_trunc=False):
        '''如果表存在,清空表数据.如果不是algtmp开头的表,需要增加强制清空参数.'''
        if not table_name.startswith('algtmp_') and not force_trunc:
            logger.info('table {table_name} is not algtmp type, truncate it anyway please use force_trunc argument.'.format(table_name=table_name))
        else:
            if self.check_table_exists(table_name):
                self.query(f'''TRUNCATE TABLE {table_name} ''')
                
    
    def execute_without_result(self, sql):
        cur = self.conn.cursor()
        try:
            cur.execute('set hive.exec.dynamic.partition.mode=nonstrict')
            cur.execute('set hive.exec.max.dynamic.partitions.pernode=10000')
            cur.execute('set hive.exec.max.dynamic.partitions=10000')
            cur.execute(sql)
        except Exception as e:
            logger.error('Error occurred when executing:')
            logger.error(sql)
            raise e
    
    
    def overwrite_data_to_partition(self, table_name, partition, data_list, batch_size=100000):
        '''
        data_list字段顺序要完全和已有表字段顺序一致:
        '''

        logger.info('saving to table "%s" partition "%s" total records: %d' % (table_name, partition, len(data_list)))
        # save with batch_size
        batch_n = 0
        names = list(data_list[0].keys())
        while batch_n * batch_size < len(data_list):
            logger.info('saved %d.' % (batch_n * batch_size))
            base_sql = 'insert overwrite table {} partition({})'.format(table_name, partition) + ' values'

            row_values = []
            for meta in data_list[batch_n * batch_size : (batch_n + 1) * batch_size]:
                values = []
                for col in names:
                    # values.append("{}".format(meta[col]))
                    values.append("'{}'".format(str(meta[col]).replace("'", "''")))
                    pass
                row_values.append('(' + ','.join(values) + ')')
            sql = base_sql + ', '.join(row_values)
            self.execute_without_result(sql)

            batch_n += 1
            pass
        logger.info('saved to table "%s" partition "%s" total records: %d' % (table_name,partition, len(data_list)))
        pass
            

if __name__ == '__main__':
    hive_client = HiveClient.instance(config.get('hive_conn_info'))
    # print(hive_client.query('show tables'))
    print(hive_client.run_sql_return_plain_json('''
        select distinct dma_id
        from qdwater_cdm.dwd_res_rltn_wateruser_dma
    '''))
    hive_client.close()