manyi/script/upgrade/executed/upgrade_source_key_20191111.py

# -*- coding: utf-8 -*-
# !/usr/bin/env python

"""

"""

import os
import sys

from bson import ObjectId
from pymongo.collection import Collection
from typing import Dict

PROJECT_ROOT = os.path.join(os.path.abspath(os.path.split(os.path.realpath(__file__))[0] + "/.."), '..')
sys.path.insert(0, PROJECT_ROOT)

os.environ.update({"DJANGO_SETTINGS_MODULE": "configs.systest"})

from script.base import init_env

init_env(interactive = False)

from apps.web.dealer.models import Dealer
from apps.web.agent.models import Agent
from apps.web.core import APP_KEY_DELIMITER, PayAppType
from apilib.monetary import RMB
from apps.web.core.payment import WithdrawGateway
from apilib.utils_mongo import BulkHandler


def refresh_update_condition(bulker, query, field_name, gateway_key, balance):
    # type: (Dict, Dict, str, str, Dict)->Dict
    if gateway_key.startswith(WithdrawGateway.LEDGER_PREFIX) or gateway_key.startswith(
            WithdrawGateway.NO_LEDGER_PREFIX):
        print('ignore. is source key = {}'.format(gateway_key))
        return

    occupant_id, appid, mchid = gateway_key.split(APP_KEY_DELIMITER)

    souce_key = APP_KEY_DELIMITER.join([WithdrawGateway.LEDGER_PREFIX, PayAppType.WECHAT, occupant_id, mchid])

    update = {}

    update['$unset'] = {
        '{filed}.{key}'.format(filed = field_name, key = gateway_key): 0
    }

    update['$inc'] = {
        '{filed}.{key}.balance'.format(filed = field_name,
                                       key = souce_key): balance.get('balance', RMB('0.00').mongo_amount),
        '{filed}.{key}.frozenBalance'.format(filed = field_name,
                                             key = souce_key): balance.get('frozenBalance', RMB('0.00').mongo_amount)
    }

    bulker.update(query_dict = query, update_dict = update)


def check_dealer():
    dealers = []

    try:
        for dealer in Dealer.objects.all():
            dealers.append(dealer)
    except Exception as e:
        print(str(e))

    new_map = {}
    for dealer in dealers:
        _balance = RMB(0)
        _frozen_balance = RMB(0)
        for gateway_key, balance in dealer.deviceBalance.iteritems():
            _balance += balance.balance
            _frozen_balance += balance.frozenBalance

        for gateway_key, balance in dealer.adBalance.iteritems():
            _balance += balance.balance
            _frozen_balance += balance.frozenBalance

        new_map[str(dealer.id)] = {
            'balance': _balance,
            'frozen_balance': _frozen_balance
        }

    from mongoengine import connection
    new_connection = connection.get_db()['Dealer_copy']
    from mongoengine.queryset import QuerySet
    new_objects = QuerySet(Dealer, new_connection)

    try:
        old_map = {}
        for dealer in new_objects.all():
            _balance = RMB(0)
            _frozen_balance = RMB(0)
            for gateway_key, balance in dealer.deviceBalance.iteritems():
                _balance += balance.balance
                _frozen_balance += balance.frozenBalance

            for gateway_key, balance in dealer.adBalance.iteritems():
                _balance += balance.balance
                _frozen_balance += balance.frozenBalance

            old_map[str(dealer.id)] = {
                'balance': _balance,
                'frozen_balance': _frozen_balance
            }

        if len(new_map.keys()) != len(old_map.keys()):
            print('dealer len is same. id = {}'.format(str(dealer.id)))

    except Exception as e:
        print(str(e))

    for _id, balance in new_map.iteritems():
        old_balance = old_map[_id]

        if balance['balance'] != old_balance['balance']:
            print('dealer balace is not equal. id = {}'.format(_id))

        if balance['frozen_balance'] != old_balance['frozen_balance']:
            print('dealer frozen_balance is not equal. id = {}'.format(_id))

        print('dealer is ok. id = {}'.format(_id))


def check_agent():
    agents = []

    try:
        for agent in Agent.objects.all():
            agents.append(agent)
    except Exception as e:
        print(str(e))

    new_map = {}
    for agent in agents:
        _balance = RMB(0)
        _frozen_balance = RMB(0)
        for gateway_key, balance in agent.deviceBalance.iteritems():
            _balance += balance.balance
            _frozen_balance += balance.frozenBalance

        for gateway_key, balance in agent.adBalance.iteritems():
            _balance += balance.balance
            _frozen_balance += balance.frozenBalance

        for gateway_key, balance in agent.withdrawBalance.iteritems():
            _balance += balance.balance
            _frozen_balance += balance.frozenBalance

        for gateway_key, balance in agent.trafficBalance.iteritems():
            _balance += balance.balance
            _frozen_balance += balance.frozenBalance

        new_map[str(agent.id)] = {
            'balance': _balance,
            'frozen_balance': _frozen_balance
        }

    from mongoengine import connection
    new_connection = connection.get_db()['Agent_copy']
    from mongoengine.queryset import QuerySet
    new_objects = QuerySet(Agent, new_connection)

    try:
        old_map = {}
        for agent in new_objects.all():
            _balance = RMB(0)
            _frozen_balance = RMB(0)
            for gateway_key, balance in agent.deviceBalance.iteritems():
                _balance += balance.balance
                _frozen_balance += balance.frozenBalance

            for gateway_key, balance in agent.adBalance.iteritems():
                _balance += balance.balance
                _frozen_balance += balance.frozenBalance

            for gateway_key, balance in agent.withdrawBalance.iteritems():
                _balance += balance.balance
                _frozen_balance += balance.frozenBalance

            for gateway_key, balance in agent.trafficBalance.iteritems():
                _balance += balance.balance
                _frozen_balance += balance.frozenBalance

            old_map[str(agent.id)] = {
                'balance': _balance,
                'frozen_balance': _frozen_balance
            }

        if len(new_map.keys()) != len(old_map.keys()):
            print('agent len is not same. id = {}'.format(str(agent.id)))

    except Exception as e:
        print(str(e))

    for _id, balance in new_map.iteritems():
        old_balance = old_map[_id]

        if balance['balance'] != old_balance['balance']:
            print('agent balace is not equal. id = {}'.format(_id))

        if balance['frozen_balance'] != old_balance['frozen_balance']:
            print('agent frozen_balance is not equal. id = {}'.format(_id))

        print('agent is ok. id = {}'.format(_id))


dealers = Dealer.get_collection().find({})
for dealer in dealers:
    bulker = BulkHandler(Dealer.get_collection())

    query = {'_id': dealer['_id']}

    if Dealer.deviceBalance.name in dealer:
        for gateway_key, balance in dealer[Dealer.deviceBalance.name].iteritems():
            refresh_update_condition(bulker, query, Dealer.deviceBalance.name, gateway_key, balance)

    if Dealer.adBalance.name in dealer:
        for gateway_key, balance in dealer[Dealer.adBalance.name].iteritems():
            refresh_update_condition(bulker, query, Dealer.adBalance.name, gateway_key, balance)

    if len(bulker.bulk._BulkOperationBuilder__bulk.ops) > 0:
        bulker.execute()

agents = Agent.get_collection().find({})
for agent in agents:
    bulker = BulkHandler(Agent.get_collection())
    query = {'_id': agent['_id']}

    if Agent.deviceBalance.name in agent:
        for gateway_key, balance in agent[Agent.deviceBalance.name].iteritems():
            refresh_update_condition(bulker, query, Agent.deviceBalance.name, gateway_key, balance)

    if Agent.adBalance.name in agent:
        for gateway_key, balance in agent[Agent.adBalance.name].iteritems():
            refresh_update_condition(bulker, query, Agent.adBalance.name, gateway_key, balance)

    if Agent.trafficBalance.name in agent:
        for gateway_key, balance in agent[Agent.trafficBalance.name].iteritems():
            refresh_update_condition(bulker, query, Agent.trafficBalance.name, gateway_key, balance)

    if Agent.withdrawBalance.name in agent:
        for gateway_key, balance in agent[Agent.withdrawBalance.name].iteritems():
            refresh_update_condition(bulker, query, Agent.withdrawBalance.name, gateway_key, balance)

    if len(bulker.bulk._BulkOperationBuilder__bulk.ops) > 0:
        bulker.execute()

check_dealer()
check_agent()