Python mongoengine 模块，connect() 实例源码

def create_app(extra_config_settings={}):
    """
    Initialize Flask applicaton
    """

    app.config.from_object('app.startup.settings')
    app.config.update(extra_config_settings)

    # Load all blueprints with their manager commands, entities and views
    from app import core

    mongoengine.connect(
        app.config['DATABASE_NAME'],
        host=app.config['DATABASE_HOST'],
        port=app.config['DATABASE_PORT'],
        username=app.config['DATABASE_USER'],
        password=app.config['DATABASE_PASSWORD']
    )

    return app

def task_before_sent_handler(sender=None, headers=None, body=None, **kwargs):
    # information about task are located in headers for task messages
    # using the task protocol version 2.
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    task_name = sender
    args = headers.get('argsrepr')
    task_id = headers.get('id')
    task_monitor_ob = task_monitor()
    task_monitor_ob.task_id = task_id
    task_monitor_ob.task_name = task_name
    task_monitor_ob.before_sent_args = args
    now = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.create_time = now
    task_monitor_ob.update_time = now
    task_monitor_ob.celery_stask_status = 0
    task_monitor_ob.save()

def mongo_connect(*args, **kwargs):
    """Connect mongoengine to mongo db. This connection is reused everywhere"""
    for _ in xrange(30):
        try:
            log.info("Attempting to connect to %s at %s...", config.MONGO_DB,
                     config.MONGO_URI)
            me.connect(db=config.MONGO_DB, host=config.MONGO_URI)
        except Exception as exc:
            log.warning("Error connecting to mongo, will retry in 1 sec: %r",
                        exc)
            time.sleep(1)
        else:
            log.info("Connected...")
            break
    else:
        log.critical("Unable to connect to %s at %s: %r", config.MONGO_DB,
                     config.MONGO_URI, exc)
        raise exc

def process_drip_campaigns_task(conf):
    # initialize mongo connection
    mongoengine.connect(
        conf.MONGODB_SETTINGS["db"],
        host=conf.MONGODB_SETTINGS["host"],
        port=conf.MONGODB_SETTINGS["port"])
    # get log directory and name and initialize logger
    log_dir = conf.LOG_SETTINGS["log_dir"]
    log_name = conf.LOG_SETTINGS["log_name"]
    file_handler = RotatingFileHandler(
        filename=os.path.join(log_dir, log_name + ".log"),
        maxBytes=200000000,  # 200MB
        backupCount=20,  # 20*200MB=4GB=nice
    )
    logger = logging.getLogger(name=log_name)
    logger.addHandler(file_handler)
    logger.setLevel(logging.INFO)
    # run the processor
    process_campaigns(logger)

def __init__(self, server_config, auto_connect=True):
        """Initialise the client

        :param server_config: The server configuration
        :param auto_connect: Whether to automatically connect
        """
        self.server_config = server_config

        if auto_connect:
            if ServerSelectionTimeoutError:
                try:
                    self.connect(server_config)
                except ServerSelectionTimeoutError as e:
                    logging.warn(e.message)
                    sys.exit()
            else:
                self.connect(server_config)

def mongodb_connect(connection, alias):
    user = connection.get('USER', None)
    password = connection.get('PASSWORD', None)
    host = connection.get('HOST', 'localhost')
    port = connection.get('PORT', 27017)
    name = connection.get('NAME', 'audit')
    replica_set = connection.get('REPLICA_SET', '')

    if isinstance(host, (list, tuple)) and isinstance(port, (list, tuple)) and len(host) == len(port):
        hosts_list = ["{}:{}".format(h, p) for h, p in zip(host, port)]
        hosts_str = ",".join(hosts_list)
        options = "?replicaSet={}".format(replica_set)
    else:
        hosts_str = "{}:{}".format(host, port)
        options = ""

    if user and password:
        uri = 'mongodb://{}:{}@{}/{}{}'.format(user, password, hosts_str, name, options)
    else:
        uri = 'mongodb://{}/{}{}'.format(hosts_str, name, options)

    try:
        mongoengine.connect(name, host=uri, alias=alias)
    except ConnectionError as e:
        logger.error('Database connection error: %s', e.message, exc_info=e)
        raise e

def get_tracer_and_connect(self):
        tracer = get_dummy_tracer()
        Pin.get_from(mongoengine.connect).clone(
            tracer=tracer).onto(mongoengine.connect)
        mongoengine.connect(port=MONGO_CONFIG['port'])

        return tracer

def get_tracer_and_connect(self):
        tracer = TestMongoEnginePatchConnectDefault.get_tracer_and_connect(self)
        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(mongoengine.connect)
        mongoengine.connect(port=MONGO_CONFIG['port'])

        return tracer

def get_tracer_and_connect(self):
        tracer = get_dummy_tracer()
        client = mongoengine.connect(port=MONGO_CONFIG['port'])
        Pin.get_from(client).clone(tracer=tracer).onto(client)

        return tracer

def get_tracer_and_connect(self):
        tracer = get_dummy_tracer()
        # Set a connect-level service, to check that we properly override it
        Pin(service='not-%s' % self.TEST_SERVICE).onto(mongoengine.connect)
        client = mongoengine.connect(port=MONGO_CONFIG['port'])
        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(client)

        return tracer

def test_patch_unpatch(self):
        tracer = get_dummy_tracer()

        # Test patch idempotence
        patch()
        patch()

        client = mongoengine.connect(port=MONGO_CONFIG['port'])
        Pin.get_from(client).clone(tracer=tracer).onto(client)

        Artist.drop_collection()
        spans = tracer.writer.pop()
        assert spans, spans
        eq_(len(spans), 1)

        # Test unpatch
        mongoengine.connection.disconnect()
        unpatch()

        mongoengine.connect(port=MONGO_CONFIG['port'])

        Artist.drop_collection()
        spans = tracer.writer.pop()
        assert not spans, spans

        # Test patch again
        patch()

        client = mongoengine.connect(port=MONGO_CONFIG['port'])
        Pin.get_from(client).clone(tracer=tracer).onto(client)

        Artist.drop_collection()
        spans = tracer.writer.pop()
        assert spans, spans
        eq_(len(spans), 1)

def unpatch():
    setattr(mongoengine, 'connect', _connect)

def connect(**kwargs):
    '''Encapsulates connecting to the configured mongodb instance
    '''
    return conn(**kwargs)

def set_up_database(self):
        # Create the plusequals-test database
        self.db = connect("dhariri-test")

def config_mongo(**kwargs):
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    DBStore._initialize()
    ProxyStore._initialize()
    log.other.info("???mongo??,pid:%s" % os.getpid())

def task_prerun_handler(task_id = None, args = None, **kwargs):
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    #information about task are located in headers for task messages
    # using the task protocol version 2.
    print("task_prerun_handler:" + str(task_id))
    task_monitor_ob = task_monitor.objects(task_id= task_id).first()
    task_monitor_ob.task_prerun_args = args
    task_monitor_ob.celery_stask_status = 1
    task_monitor_ob.update_time = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.save()

def task_success_handler(sender=None, headers=None, body=None, **kwargs):
    # information about task are located in headers for task messages
    # using the task protocol version 2.
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    task_id = sender.request.get('id')
    print("task_success_handler:" + str(task_id))
    task_monitor_ob = task_monitor.objects(task_id= task_id).first()
    task_monitor_ob.celery_stask_status = 5
    task_monitor_ob.update_time = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.save()

def task_failure_handler(sender=None, headers=None, body=None, **kwargs):
    # information about task are located in headers for task messages
    # using the task protocol version 2.
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    task_id = sender.request.get('id')
    task_monitor_ob = task_monitor.objects(task_id= task_id).first()
    task_monitor_ob.celery_stask_status = 6
    task_monitor_ob.update_time = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.save()

# ?????

def task_prerun_handler(task_id = None, args = None, **kwargs):
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    #information about task are located in headers for task messages
    # using the task protocol version 2.
    print("task_prerun_handler:" + str(task_id))
    task_monitor_ob = task_monitor.objects(task_id= task_id).first()
    task_monitor_ob.task_prerun_args = args
    task_monitor_ob.celery_stask_status = 1
    task_monitor_ob.update_time = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.save()

def task_success_handler(sender=None, headers=None, body=None, **kwargs):
    # information about task are located in headers for task messages
    # using the task protocol version 2.
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    task_id = sender.request.get('id')
    print("task_success_handler:" + str(task_id))
    task_monitor_ob = task_monitor.objects(task_id= task_id).first()
    task_monitor_ob.celery_stask_status = 5
    task_monitor_ob.update_time = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.save()

def task_failure_handler(sender=None, headers=None, body=None, **kwargs):
    # information about task are located in headers for task messages
    # using the task protocol version 2.
    mongoengine.connect(**celery_config.mongoengine_SETTINGS)
    task_id = sender.request.get('id')
    task_monitor_ob = task_monitor.objects(task_id= task_id).first()
    task_monitor_ob.celery_stask_status = 6
    task_monitor_ob.update_time = datetime.now(tz = pytz.timezone('Asia/Shanghai'))
    task_monitor_ob.save()

def connectMongoEngine(pmcol, conn_uri=None):
    if pymongo.version_tuple[0] == 2:     #really? REALLY?
        #host = pmcol.database.connection.HOST
        #port = pmcol.database.connection.PORT
        host = pmcol.database.connection.host
        port = pmcol.database.connection.port
    else:
        host = pmcol.database.client.HOST
        port = pmcol.database.client.PORT
    # Can just use the connection uri, which has credentials
    if conn_uri:
        return meng.connect(pmcol.database.name, host=conn_uri)
    return meng.connect(pmcol.database.name, host=host, port=port)

def __connect_db(self):
        try:
            app.logger.info(("Checking the connectivity to the database(%s)...." % self.__database_ip))
            db =  connect(self.__database_name, host=self.__database_ip, port=self.__port, serverSelectionTimeoutMS=3000)
            db.server_info()
            app.logger.info("...OK")
            return db
        except Exception as e:
            app.logger.error("Unable to connect to %s" % self.__database_ip)
            sys.exit(1)

def setUpClass(cls):
        """Setup database connection."""
        cls.db = db.connect("goodjson_test", host='mongomock://localhost')

def _connect(conn_settings):
    """Given a dict of connection settings, create a connection to
    MongoDB by calling mongoengine.connect and return its result.
    """
    db_name = conn_settings.pop('name')
    return mongoengine.connect(db_name, **conn_settings)

def open_spider(self, _):
        logger.info("Connecting to %s", self.mongo_url)
        self.client = connect(host=self.mongo_url)
        logger.debug("Connected")

def cli(args):
    try:
        opts, args = getopt.getopt(args, "i:", ["id=", ])
    except getopt.GetoptError:
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print('A simple script for loading data from PostgreSQL to Mongo')
            sys.exit()
        elif opt in ("-i", "--id"):
            row_id = arg

            # Connection to the PostgreSQL, to be defined in the Airflow UI
            pg_hook = PostgresHook(postgres_conn_id="postgres_data")

            # Retrieve the data stored in PostgreSQL
            pg_command = """SELECT * FROM dag_dag WHERE id = %s"""
            data = pg_hook.get_records(pg_command, parameters=[row_id])

            # Connect to Mongo databases in the Docker compose
            mongoengine.connect(db="dags", host="mongo:27017", alias="default")

            # Search for existing documents with the same dag_name
            dags_docs = DAG_Description.objects(dag_name="hello_world_template")

            # logging.info(type(dags_docs))
            # logging.info(dags_docs)

            # If there are no documents
            if len(dags_docs) == 0:
                # Setup a new document for storing the data
                logging.info("Creating a new Mongo document for %s", row_id)
                dag_document = DAG_Description(
                    dag_name=row_id,
                    raw_data=data[0][1],
                    clean_data=data[0][2],
                    vis_type=data[0][3],
                    vis_title=data[0][4],
                    vis_text=data[0][5],
                    vis_footer=data[0][6],
                    created_at=data[0][7],
                    updated_at=data[0][8])
                # Save the document
                dag_document.save()
            # If there are more than one documents, get the first one
            elif len(dags_docs) > 1:
                logging.info("Updating the first Mongo document found for %s", row_id)
                dag_document = dags_docs.first()
                dag_document.update(raw_data=data[0][1], clean_data=data[0][2], updated_at=datetime.now)
            # If there is only one document
            elif len(dags_docs) == 1:
                logging.info("Updating the Mongo document found for %s", row_id)
                dag_document = dags_docs.first()
                dag_document.update(raw_data=data[0][1], clean_data=data[0][2], updated_at=datetime.now)

            # Return the success message
            logging.info("Data exported from PostgreSQL to Mongo successfully.")

def connect(self, server_config):
        """Connect using the configuration given

        :param server_config: The server configuration
        """
        if 'connection_string' in server_config:
            self.client = pymongo.MongoClient(
                    server_config['connection_string'])
            self.db = self.client[server_config['db']]
        else:
            self.client = pymongo.MongoClient(
                server_config['host'],
                server_config['port'],
                tz_aware=self.get_config_value('tz_aware', True))

            self.db = self.client[server_config['db']]

        if ('authentication_database' in server_config and
                server_config['authentication_database']):
            self.db.authenticate(
                    server_config['username'], server_config['password'],
                    source=server_config['authentication_database'])
        else:
            if 'username' in server_config:
                if 'password' in server_config:
                    self.db.authenticate(server_config['username'],
                                         server_config['password'])
                else:
                    self.db.authenticate(server_config['username'])

        # Mongo Engine connection
        d = dict((k, v) for k, v in server_config.items()
                 if k not in ['modalities', 'summaries'])
        if 'authentication_database' in d:
            d['authentication_source'] = d['authentication_database']
            del d['authentication_database']

        self.session = connect(alias="hyperstream", **d)

        # TODO: This sets the default connection of mongoengine, but seems to be a bit of a hack
        if "default" not in connection._connections:
            connection._connections["default"] = connection._connections["hyperstream"]
            connection._connection_settings["default"] = connection._connection_settings["hyperstream"]