BIDI_EXTRACTION_AI/BiddingKG/dl/infra/mongo.py

# -*- coding: utf-8 -*-
"""MongoDB 客户端工厂。

替代原 entityLink/mongotest.py 中硬编码的连接字符串。
URI 走环境变量 BIDDINGKG_MONGO_URI，不写入代码。
"""
from __future__ import absolute_import

import logging

from . import config

__all__ = ["get_mongo_client", "get_mongo_db"]

logger = logging.getLogger(__name__)


def get_mongo_client():
    """返回 pymongo MongoClient。未配置 URI 时抛出明确错误。"""
    try:
        import pymongo
    except ImportError:  # pragma: no cover
        raise ImportError("pymongo is required: pip install pymongo")

    uri = config.mongo_uri()
    if not uri:
        raise RuntimeError(
            "MongoDB URI not configured. Set BIDDINGKG_MONGO_URI env var or "
            "config/.env."
        )
    return pymongo.MongoClient(uri)


def get_mongo_db(dbname="bxkc"):
    """返回指定 db。"""
    client = get_mongo_client()
    return client[dbname]