tfse-model-api-v0.2/Report/db.py

"""
Mongo_tfse
"""

import pymongo
from bson import ObjectId
from gridfs import GridFS

from Report.path_tool import get_gen_report_path

DB_HOST = '116.63.130.34'
DB_USER = 'root'
DB_PASS = 'sromitdTW569kC#M'
DB_PORT = 27018

client = pymongo.MongoClient('mongodb://{}:{}@{}:{}'.format(DB_USER, DB_PASS, DB_HOST, DB_PORT))


def upload_pdf(name):
    """
    根据名称保存该企业报告
    """
    path = get_gen_report_path(name=name)
    database = '企业'
    collection = 'PDF'
    fs = GridFS(client[database], collection)
    content = open(path, 'rb').read()
    pdf_id = fs.put(content, content_type='application/pdf', filename=name)

    return str(pdf_id)


def download_pdf(pdf_id):
    """
    根据pdf_id查询pdf文件
    """
    database = '企业'
    collection = 'PDF'
    fs = GridFS(client[database], collection)
    data = fs.get(ObjectId(pdf_id)).read()

    return data


def find_tfse_db(**kwargs):
    """
    查询数据库记录
    Parameters:
        kwargs: 查询条件
    Returns:
        data: tfse数据
    """
    collection = client[kwargs['db']][kwargs['col']]
    result = collection.find(kwargs['query'], {'_id': False})
    data = list(result)
    return data


if __name__ == '__main__':
    pid = upload_pdf('远东资信评估有限公司')
    file = download_pdf(pid)