poyee-data-warehouse/dw_base/scheduler/mg_company_alias_init.py

# mongo公司别名表索引初始化
import sys
import re
import os

abspath = os.path.abspath(__file__)
root_path = re.sub(r"tendata-warehouse.*", "tendata-warehouse", abspath)
sys.path.append(root_path)
from dw_base.utils.config_utils import parse_args
from configparser import ConfigParser
from pymongo import MongoClient


def get_mongo_client(conf_path):
    config_parser = ConfigParser()
    config_parser.read(root_path + conf_path)
    print('conf_path:', root_path + conf_path)
    url = config_parser.get('base', 'address')
    return MongoClient(url)


def create_index(client, tbl_name):
    collection = client['tendata_corp'][tbl_name]
    collection.create_index([("tid", 1)], unique=True)
    collection.create_index([("qybzmc", 1)])


def check_index(client, tbl_name):
    collection = client['tendata_corp'][tbl_name]
    index_info = str(collection.index_information())
    cnt = collection.count()
    print(f'{tbl_name} count: {cnt}')
    if "'key': [('tid', 1)], 'unique': True" in index_info and "'key': [('qybzmc', 1)]" in index_info:
        return True
    else:
        return False


if __name__ == '__main__':
    ent_mg_conf = '/../datasource/mongo/mongo-ent-prod-alias-rw.ini'
    CONFIG, _ = parse_args(sys.argv[1:])
    country = CONFIG.get('country')
    tbl_name = f'company_alias_{country}'
    client = get_mongo_client(ent_mg_conf)
    create_index(client, tbl_name)
    if check_index(client, tbl_name):
        print(f'{tbl_name}: index ok')
    else:
        print(f'{tbl_name}: index not ok,plz check!')