Загрузка данных


import socket
import sys
import io
import json
import traceback
from collections import OrderedDict
from datetime import datetime, timedelta
from dateutil.parser import isoparse
import time
import builtins
import requests
import requests_cache
from lxml.builder import E
from lxml import etree
from retry import retry

XML_DEFAULT_BATCH_SIZE  = 100000
CSV_DEFAULT_BATCH_SIZE  = 300000
DATA_SOURCE_TYPE_URL    = 'url'
DATA_SOURCE_TYPE_FILE   = 'file'
DATA_SOURCE_TYPE_ALL    = 'all'
DATA_SOURCE_TYPE_LIST   = (DATA_SOURCE_TYPE_URL, DATA_SOURCE_TYPE_FILE, DATA_SOURCE_TYPE_ALL)
IMPORT_TOKEN = 'EXTERNAL_SERVICE'


def get_etl_descriptions(application_name, etl_desc_url):
    log(f'Getting ETL descriptions of {application_name}')
    with requests_cache.disabled():
        responce = requests.get(etl_desc_url)
    assert responce.status_code==200, \
           f'Error getting ETLs descriptions, responce code {responce.status_code}, {responce.text}'
    etl_descriptions = []
    for etl in json.loads(json.loads(responce.text)['RETURN']):
        if etl['application_name']==application_name:
            etl_descriptions.append(etl)
    log(f'Found {len(etl_descriptions)} descriptions')
    return etl_descriptions

def csv_to_dictlist(raw_data, source_description):
    import csv
    log('Parsing csv')
    result = []
    header_list = []
    row_id = 0
    start_marker = source_description.get('start_marker', None)
    end_marker = source_description.get('end_marker', None)
    in_data = start_marker is None
    reader = csv.reader(raw_data.splitlines(), delimiter=DELIMITER)
    for row in reader:
        if in_data:
            # stop parsing csv if end_marker is reached
            if end_marker is not None and end_marker in row: break
            # skip empty rows
            if not any(value!='' for value in row): continue
            # read header, rename empty header columns to columnx, x - column number
            if not header_list: 
                header_list = [value.strip() if value!='' else f'column{id}' for id, value in enumerate(row)]
                continue
            row_dict = OrderedDict()
            for col, value in enumerate(row): 
                row_dict[header_list[col]] = value
            if row_id%CSV_DEFAULT_BATCH_SIZE==0 and result:
                yield result
                result = []
            result.append(row_dict)
            row_id += 1
        else:
            in_data = start_marker in row
    if result: yield result

def fstring_eval(fstring):
    return eval(fstring) if fstring.strip()[:2]=="f'" else fstring

def process_raw_dictlist(dl_raw_data, source_description, source_id, columns_defaults):

    def parse_value(value, type):
        try:
            value = value.strip()
            # CWG Solar Observations may return NULL as value; Weather historical - nan; Hydro - n/a
            if value.upper() in ('NULL', 'NAN', 'N/A'): return ''
            if type!='date': return value
            return '' if value=='' else DATE_PARSER(value).strftime(XML_DATE_FMT)
        except:
            log(f'Error parsing [{value}] date value {traceback.format_exception(*sys.exc_info())}')
            return ''

    log('Converting raw data')
    # python3 exec() is buggy with locals()
    # so we have to declare result_dict as global for correct post processing
    # https://bugs.python.org/issue4831 
    global result_dict
    result_dict = []
    unpivot_from_column = None
    xml_chunk = source_id if source_description.get('split_xml_by_url', 'N')=='Y' else -1
    for raw_row in dl_raw_data:
        row_dict = OrderedDict()
        row_dict['xml_chunk'] = xml_chunk
        for attribute in source_description['attributes']:
            origin = attribute.get('origin', None)
            if attribute.get('default', None) is not None:
                columns_defaults[attribute['destination']] = attribute['default']
            if origin:
                origin = fstring_eval(origin)
                row_dict[attribute['destination']] = parse_value(raw_row.get(origin, ''), attribute['type'])
                continue
            value = attribute.get('value', None)
            if value:
                value = fstring_eval(value)
                row_dict[attribute['destination']] = parse_value(value, attribute['type'])
                continue
            if unpivot_from_column is None:
                unpivot_from_column       = attribute.get('unpivot_from_column', None)
                unpivot_destination       = attribute.get('destination', None)
                unpivot_type              = attribute.get('type', None)
                unpivot_value_destination = attribute.get('value_destination', None)
                unpivot_value_type        = attribute.get('value_type', None)
        # check the need of unpivot
        if unpivot_from_column is not None:
            for col, item in enumerate(raw_row.items()):
                # ignore empty pivoting header elements (check item[0] is not like columnx)
                if item[0][:6]!='column' and col >= unpivot_from_column:
                    result_dict.append(
                        {
                            **row_dict,
                            unpivot_destination: parse_value(item[0], unpivot_type),
                            unpivot_value_destination: parse_value(item[1], unpivot_value_type)
                        }
                    )
        else:
           result_dict.append(row_dict)
    # postprocessing
    if source_description.get('post_processor', None):
        log('Post processing data')
        exec('\n'.join(source_description['post_processor']), globals())

    return result_dict

def get_source_iterator(source_description):
    try:
        source_iterator_code = '\n'.join(source_description.get('source_iterator', ['def source_iterator():', '  yield None']))
        exec(source_iterator_code, globals())
        return source_iterator
    except:
        log(f'Error compiling source iterator code {source_iterator_code} {traceback.format_exception(*sys.exc_info())}')
        raise

def wait(source_description):
    request_delay = source_description.get('request_delay', 0)
    if request_delay!=0:
        log(f'Waiting {request_delay} seconds before next request')
        time.sleep(request_delay)

def file_requestor(file):
    with io.open(file, mode='r', encoding='utf-8') as file:
        return file.read()

def get_processed_data(etl_description, api_requestor, data_source_type, columns_defaults):
    processed_data = []
    # declare raw_data and source_description as global for correct pre processing
    global raw_data
    global source_description
    global do_request
    for source_description in etl_description['sources']:
        if source_description['source_type']!=data_source_type and data_source_type!=DATA_SOURCE_TYPE_ALL: continue
        builtins.DATE_PARSER = isoparse
        builtins.source_description = source_description
        builtins.etl_description = etl_description
        for source_id, ui in enumerate(get_source_iterator(source_description)()):
            try:
                source = fstring_eval(source_description['source'])
                #log (f'Pre_request')
                if source_description.get('pre_request', None):
# Проверяем
# Если в конфиге есть элемент "pre_request" и питоновский код в нём вернет False, то прекращаем обработку текущего конфига и переходим к следующей.
                    do_request = True
                    log('PreRequest raw data')

                    exec('\n'.join(source_description['pre_request']), globals())
                    if do_request == False:
                        log ('pre_request returns FALSE. Do next itteration of source.')
                        continue

                log(f'Processing source {source}')
                raw_data = api_requestor(url=source) if source_description['source_type']==DATA_SOURCE_TYPE_URL else file_requestor(file=source)

# preprocessing raw data
                if source_description.get('pre_processor', None):
                    log('Preprocessing raw data')
                    exec('\n'.join(source_description['pre_processor']), globals())

                for dl_raw_data in csv_to_dictlist(raw_data, source_description):
                    processed_data += process_raw_dictlist(dl_raw_data, source_description, source_id, columns_defaults)
                    if len(processed_data) >= CSV_DEFAULT_BATCH_SIZE:
                        yield processed_data
                        processed_data = []
                wait(source_description)
            except:
                log(f'Error processing data {traceback.format_exception(*sys.exc_info())}')
    if processed_data: yield processed_data

# implemented xml batching because:
# 1. lxml.etree.tostring can't serialize ~300Mb xml - MemoryError
# 2. ~300Mb xml serialization by io.StringIO() works fine, but requests.post fails with MemoryError
def generate_xml(processed_data, etl_description, columns_defaults):

    import hashlib
    def row_hash(row):
        return hashlib.sha256('{|}'.join([value for key, value in row.items() if value!='' and key in hashed_column_list]).encode()).hexdigest()

    log('Generating xml')
    xml = []
    db_export_column_list = []
    hashed_column_list = []
    hash_needed = etl_description['merge_list'] == 'hash' if 'merge_list' in etl_description else False
    for attr in etl_description['sources'][0]['attributes']:
        if attr.get('db_export', 'Y')=='Y':
            db_export_column_list.append(attr['destination'])
        if attr.get('db_export', 'Y')=='Y' and attr.get('hash_exclude', 'N')=='N':
            hashed_column_list.append(attr['destination'])
    xml_batch_size = etl_description.get('xml_batch_size', XML_DEFAULT_BATCH_SIZE)
    xml_chunk = None
    for row_id, row in enumerate(processed_data):
        if (row_id%xml_batch_size==0 or row['xml_chunk']!=xml_chunk) and xml:
            yield etree.tostring(
                E.lng(E.type(etl_description['import_name']), E.rows(*xml)),
                encoding='unicode',
                pretty_print=True)
            xml = []
        xml_chunk = row['xml_chunk']
        xml_row = list()
        for key, value in row.items():
            if value == '': value = columns_defaults.get(key, '')
            if value == '' or key not in db_export_column_list: continue
            xml_row.append(E(key, value))
        if hash_needed: xml_row.append({'hash':row_hash(row)})
        xml.append(E.row(*xml_row))
    if xml:
        yield etree.tostring(
            E.lng(E.type(etl_description['import_name']), E.rows(*xml)),
            encoding='unicode',
            pretty_print=True)

def process_etls(application_name,
                 etl_desc_url,
                 data_type,
                 api_requestor,
                 erouter_url_list,
                 data_source_type=DATA_SOURCE_TYPE_ALL,
                 after_xml_wait=0,
                 import_group=None):

    def upload_xml(data_type):
        for url in erouter_url_list:
            try:
                log(f'Sending csv#{csv_id} xml#{xml_id} batch to {url}')
                upload_to_warehouse(
                    url,
                    application_name,
                    data_type,
                    f'{etl_description["data_type"]}_csv{csv_id}_xml{xml_id}_batch_{datetime.now().strftime("%Y%m%d_%H%M%S")}.xml',
                    xml
                )
                if after_xml_wait:
                    time.sleep(after_xml_wait)
            except:
                log(f'Error occured {traceback.format_exception(*sys.exc_info())}')

    log(f'------------ Processing ETLs of {application_name} ------------')
    for etl_description in get_etl_descriptions(application_name, etl_desc_url):
        if data_type and data_type!=etl_description['data_type']:
            continue
        if import_group and import_group!=etl_description['import_group']:
            continue
        log(f'------ Processing {etl_description["data_type"]} ------')
        columns_defaults = dict()
        for csv_id, processed_data in enumerate(get_processed_data(etl_description, api_requestor, data_source_type, columns_defaults)):
            for xml_id, xml in enumerate(generate_xml(processed_data, etl_description, columns_defaults)):
                upload_xml(etl_description['data_type'])


@retry(tries=10, delay=1)
def iLog(log_str='', add_date=True, console_print=True, log_file=None):
    """Log message to file and console"""
    log_str_tmp = f'{datetime.now().strftime("%Y/%m/%d %H:%M:%S")} {log_str}' if add_date else log_str
    with io.open(log_file, mode='a', encoding='utf-8') as file:
        file.write(f'{log_str_tmp}\n')
    if console_print:
        print(log_str_tmp)

@retry(tries=10, delay=5)
def upload_to_warehouse(url, application_name, data_type, file_name, content):
    """Load content to ETL loader"""
    data = {'application' : application_name,
            'data_type'   : data_type,
            'fileName'    : file_name,
            'from'        : socket.gethostname(),
            'content'     : content,
            'token'       : IMPORT_TOKEN}
    with requests_cache.disabled():
        responce = requests.post(url, data=data, verify=False)
        assert responce.status_code==200, \
               f'{url} responce code {responce.status_code}, message {responce.text}'